deerflow2/frontend/src/components/workspace/messages/markdown-content.tsx

"use client";

import { CheckIcon, CopyIcon } from "lucide-react";
import { useCallback, useMemo, useState, type MouseEvent } from "react";
import type {
  AnchorHTMLAttributes,
  ComponentPropsWithoutRef,
  ReactNode,
} from "react";

import {
  MessageResponse,
  type MessageResponseProps,
} from "@/components/ai-elements/message";
import { useI18n } from "@/core/i18n/hooks";
import { streamdownPlugins } from "@/core/streamdown";
import { cn, copyToClipboard } from "@/lib/utils";

import { CitationLink } from "../citations/citation-link";

function isExternalUrl(href: string | undefined): boolean {
  return !!href && /^https?:\/\//.test(href);
}

export type MarkdownContentProps = {
  content: string;
  isLoading: boolean;
  rehypePlugins: MessageResponseProps["rehypePlugins"];
  className?: string;
  remarkPlugins?: MessageResponseProps["remarkPlugins"];
  components?: MessageResponseProps["components"];
};

type TableData = {
  headers: string[];
  rows: string[][];
};

function parseTableData(table: HTMLTableElement): TableData {
  const headers = Array.from(table.querySelectorAll("thead th")).map((cell) =>
    (cell.textContent ?? "").trim(),
  );
  const rows = Array.from(table.querySelectorAll("tbody tr")).map((row) =>
    Array.from(row.querySelectorAll("td")).map((cell) =>
      (cell.textContent ?? "").trim(),
    ),
  );
  return { headers, rows };
}

function toMarkdownTable(data: TableData): string {
  if (data.headers.length === 0) return "";
  const headerLine = `| ${data.headers.join(" | ")} |`;
  const dividerLine = `| ${data.headers.map(() => "---").join(" | ")} |`;
  const rowLines = data.rows.map((row) => `| ${row.join(" | ")} |`);
  return [headerLine, dividerLine, ...rowLines].join("\n");
}

function MarkdownTable({
  className,
  children,
  isLoading,
  copyLabel,
  ...props
}: ComponentPropsWithoutRef<"table"> & {
  isLoading: boolean;
  copyLabel: string;
}) {
  const [copied, setCopied] = useState(false);

  const handleCopy = useCallback(
    async (event: MouseEvent<HTMLButtonElement>) => {
      const wrapper = event.currentTarget.closest(
        '[data-streamdown="table-wrapper"]',
      );
      const table = wrapper?.querySelector("table");
      if (!(table instanceof HTMLTableElement)) return;

      const markdown = toMarkdownTable(parseTableData(table));
      if (!markdown) return;

      try {
        await copyToClipboard(markdown);
        setCopied(true);
        window.setTimeout(() => setCopied(false), 2000);
      } catch {
        // no-op
      }
    },
    [],
  );

  return (
    <div
      className="my-4 flex flex-col space-y-2"
      data-streamdown="table-wrapper"
    >
      <div className="flex items-center justify-end gap-1">
        <button
          className="text-muted-foreground hover:text-foreground cursor-pointer p-1 transition-all disabled:cursor-not-allowed disabled:opacity-50"
          disabled={isLoading}
          onClick={handleCopy}
          title={copyLabel}
          type="button"
        >
          {copied ? <CheckIcon size={14} /> : <CopyIcon size={14} />}
        </button>
      </div>
      <div className="overflow-x-auto">
        <table
          className={cn(
            "border-border w-full border-collapse border",
            className,
          )}
          data-streamdown="table"
          {...props}
        >
          {children}
        </table>
      </div>
    </div>
  );
}

/** Renders markdown content. */
export function MarkdownContent({
  content,
  isLoading,
  rehypePlugins,
  className,
  remarkPlugins = streamdownPlugins.remarkPlugins,
  components: componentsFromProps,
}: MarkdownContentProps) {
  const { t } = useI18n();

  const components = useMemo(() => {
    return {
      a: (props: AnchorHTMLAttributes<HTMLAnchorElement>) => {
        if (typeof props.children === "string") {
          const match = /^citation:(.+)$/.exec(props.children);
          if (match) {
            const [, text] = match;
            return <CitationLink {...props}>{text}</CitationLink>;
          }
        }
        const { className, target, rel, ...rest } = props;
        const external = isExternalUrl(props.href);
        return (
          <a
            {...rest}
            className={cn(
              "text-primary decoration-primary/30 hover:decoration-primary/60 underline underline-offset-2 transition-colors",
              className,
            )}
            target={target ?? (external ? "_blank" : undefined)}
            rel={rel ?? (external ? "noopener noreferrer" : undefined)}
          />
        );
      },
      table: ({
        children,
        className,
        ...props
      }: ComponentPropsWithoutRef<"table"> & { children?: ReactNode }) => (
        <MarkdownTable
          className={className}
          copyLabel={t.clipboard.copyToClipboard}
          isLoading={isLoading}
          {...props}
        >
          {children}
        </MarkdownTable>
      ),
      ...componentsFromProps,
    };
  }, [componentsFromProps, isLoading, t.clipboard.copyToClipboard]);

  if (!content) return null;

  return (
    <MessageResponse
      className={className}
      isAnimating={isLoading}
      controls={{ table: false }}
      parseIncompleteMarkdown={!isLoading}
      remarkPlugins={remarkPlugins}
      rehypePlugins={rehypePlugins}
      components={components}
    >
      {content}
    </MessageResponse>
  );
}