deerflow-code/frontend-web/src/core/utils/mermaid.ts

import { getFileName } from "./files";

let mermaidInitialized = false;

/**
 * Keep sandbox diagrams on the SVG-only Mermaid rendering path.
 *
 * Mermaid's default HTML labels are implemented with SVG `foreignObject`.
 * Although that looks a little richer in current browsers, many Windows 7
 * Chromium shells either do not support it or calculate its dimensions
 * incorrectly, leaving nodes blank or overlapping. Native SVG text works in
 * the Chrome 100 baseline used by this application and in older Win7 browser
 * shells as well.
 *
 * `htmlLabels` is marked secure so a diagram's front matter/directive cannot
 * switch this preview back to `foreignObject` rendering.
 */
export const SANDBOX_MERMAID_CONFIG = {
  startOnLoad: false,
  securityLevel: "strict" as const,
  htmlLabels: false,
  fontFamily: '"Microsoft YaHei", "Microsoft YaHei UI", SimSun, Arial, sans-serif',
  flowchart: {
    // A simple path is more reliable than the spline renderer in older
    // Chromium SVG implementations and is easier to read at small sizes.
    curve: "linear" as const,
    useMaxWidth: false,
  },
  secure: [
    "secure",
    "securityLevel",
    "startOnLoad",
    "maxTextSize",
    "suppressErrorRendering",
    "maxEdges",
    "htmlLabels",
  ],
};

/**
 * Models often emit almost-Mermaid that the parser rejects:
 * - `A --> B|label| C` instead of `A --> B -->|label| C`
 * - `class,E,F loss;` (comma after `class`) instead of `class E,F loss;`
 * - `class A,B input,module;` (one class name per node) instead of
 *   separate `class` statements
 *
 * Already-valid `-->|label|` links and single-class assignments are kept.
 */
export function sanitizeMermaidSource(source: string): string {
  return sanitizeMermaidClassStatements(
    source
      .replace(
        /(^|[\s;])([A-Za-z][\w-]*|[\u4e00-\u9fff]+)\|([^|\n]+)\|/gm,
        "$1$2 -->|$3|",
      )
      .replace(/(\])\|([^|\n]+)\|/g, "$1 -->|$2|"),
  );
}

const CLASS_NODE_LIST = /[A-Za-z][\w-]*(?:\s*,\s*[A-Za-z][\w-]*)*/;
const CLASS_NAME_LIST = /[A-Za-z][\w-]*(?:\s*,\s*[A-Za-z][\w-]*)*/;

function sanitizeMermaidClassStatements(source: string): string {
  return source.replace(
    /^([ \t]*)class(?!Def\b)[ \t]*,?[ \t]*(\S[^\n]*)$/gm,
    (full, indent: string, rest: string) => {
      const trimmed = rest.replace(/;+\s*$/, "").trim();
      const matched = trimmed.match(
        new RegExp(`^(${CLASS_NODE_LIST.source})\\s+(${CLASS_NAME_LIST.source})$`),
      );
      if (!matched) {
        return full;
      }
      const nodeIds = matched[1].split(/\s*,\s*/).filter(Boolean);
      const classNames = matched[2].split(/\s*,\s*/).filter(Boolean);
      if (nodeIds.length === 0 || classNames.length === 0) {
        return full;
      }
      if (classNames.length === 1) {
        return `${indent}class ${nodeIds.join(",")} ${classNames[0]};`;
      }
      const groups: Array<{ name: string; ids: string[] }> = [];
      for (let index = 0; index < nodeIds.length; index += 1) {
        const name = classNames[index] ?? classNames[classNames.length - 1];
        const last = groups[groups.length - 1];
        if (last && last.name === name) {
          last.ids.push(nodeIds[index]);
        } else {
          groups.push({ name, ids: [nodeIds[index]] });
        }
      }
      return groups
        .map((group) => `${indent}class ${group.ids.join(",")} ${group.name};`)
        .join("\n");
    },
  );
}

export function isMermaidFile(filepath: string): boolean {
  return /\.mmd$/i.test(filepath);
}

export function getMermaidImageFileName(filepath: string, extension = "svg"): string {
  return getFileName(filepath).replace(/\.mmd$/i, `.${extension}`);
}

async function getMermaidModule() {
  const mod = await import("mermaid");
  const mermaid = mod.default;
  if (!mermaidInitialized) {
    mermaid.initialize(SANDBOX_MERMAID_CONFIG);
    mermaidInitialized = true;
  }
  return mermaid;
}

export async function renderMermaidToSvg(source: string): Promise<string> {
  const text = sanitizeMermaidSource(source).trim();
  if (!text) {
    throw new Error("Mermaid 内容为空，无法生成图片");
  }

  const mermaid = await getMermaidModule();
  const renderId = `mermaid-preview-${Math.random().toString(36).slice(2, 10)}`;
  const host = document.createElement("div");
  host.style.position = "fixed";
  host.style.left = "-10000px";
  host.style.top = "-10000px";
  host.style.pointerEvents = "none";
  host.style.opacity = "0";
  document.body.appendChild(host);

  try {
    const { svg, bindFunctions } = await mermaid.render(renderId, text);
    host.innerHTML = svg;
    bindFunctions?.(host);
    return host.innerHTML;
  } finally {
    host.remove();
  }
}

export async function openMermaidSvgInNewTab(source: string): Promise<void> {
  const svg = await renderMermaidToSvg(source);
  const blob = new Blob([svg], { type: "image/svg+xml;charset=utf-8" });
  const objectUrl = URL.createObjectURL(blob);
  const win = window.open(objectUrl, "_blank", "noopener,noreferrer");
  if (win) {
    win.opener = null;
  }
  window.setTimeout(() => URL.revokeObjectURL(objectUrl), 120_000);
}

export async function downloadMermaidSvgFile(source: string, filepath: string): Promise<void> {
  const svg = await renderMermaidToSvg(source);
  const blob = new Blob([svg], { type: "image/svg+xml;charset=utf-8" });
  const objectUrl = URL.createObjectURL(blob);
  const a = document.createElement("a");
  a.href = objectUrl;
  a.download = getMermaidImageFileName(filepath);
  a.rel = "noopener";
  document.body.appendChild(a);
  a.click();
  a.remove();
  window.setTimeout(() => URL.revokeObjectURL(objectUrl), 30_000);
}