From e70d2e2465745e6d6f7ff3b96594ff9ea28b00e1 Mon Sep 17 00:00:00 2001
From: Ryan Heida <heidarih313@gmail.com>
Date: Tue, 23 Jun 2026 12:31:17 +0330
Subject: [PATCH] feat(markdown): add Mermaid diagram rendering support (#1497)

- Add MermaidDiagram component for rendering mermaid code blocks
- Integrate mermaid support into MarkdownCodeBlock
- Improve fallback handling for invalid diagrams
- Optimize mermaid initialization and theme handling
---
 .../assistant-ui/markdown-code-block.tsx      |  10 +-
 .../assistant-ui/mermaid-diagram.tsx          | 134 ++++++++++++++++++
 2 files changed, 143 insertions(+), 1 deletion(-)
 create mode 100644 surfsense_web/components/assistant-ui/mermaid-diagram.tsx

diff --git a/surfsense_web/components/assistant-ui/markdown-code-block.tsx b/surfsense_web/components/assistant-ui/markdown-code-block.tsx
index e6c735d1e..88b0916b8 100644
--- a/surfsense_web/components/assistant-ui/markdown-code-block.tsx
+++ b/surfsense_web/components/assistant-ui/markdown-code-block.tsx
@@ -6,6 +6,7 @@ import { memo, useEffect, useState } from "react";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { materialDark, materialLight } from "react-syntax-highlighter/dist/esm/styles/prism";
 
+import { MermaidDiagram } from "@/components/assistant-ui/mermaid-diagram";
 import { Button } from "@/components/ui/button";
 import { cn, copyToClipboard } from "@/lib/utils";
 
@@ -40,6 +41,7 @@ function MarkdownCodeBlockComponent({
 	isDarkMode,
 }: MarkdownCodeBlockProps) {
 	const [hasCopied, setHasCopied] = useState(false);
+	const normalizedLanguage = language.toLowerCase();
 
 	useEffect(() => {
 		if (!hasCopied) return;
@@ -47,7 +49,7 @@ function MarkdownCodeBlockComponent({
 		return () => clearTimeout(timer);
 	}, [hasCopied]);
 
-	return (
+	const codeBlock = (
 		<div className="mt-4 overflow-hidden rounded-md bg-accent">
 			<div className="flex items-center justify-between gap-4 px-4 py-2 font-semibold text-muted-foreground text-sm">
 				<span className="lowercase text-xs">{language}</span>
@@ -78,6 +80,12 @@ function MarkdownCodeBlockComponent({
 			</SyntaxHighlighter>
 		</div>
 	);
+
+	if (normalizedLanguage === "mermaid") {
+		return <MermaidDiagram source={codeText} isDarkMode={isDarkMode} fallback={codeBlock} />;
+	}
+
+	return codeBlock;
 }
 
 export const MarkdownCodeBlock = memo(MarkdownCodeBlockComponent);
diff --git a/surfsense_web/components/assistant-ui/mermaid-diagram.tsx b/surfsense_web/components/assistant-ui/mermaid-diagram.tsx
new file mode 100644
index 000000000..50f1dc6de
--- /dev/null
+++ b/surfsense_web/components/assistant-ui/mermaid-diagram.tsx
@@ -0,0 +1,134 @@
+"use client";
+
+import { CheckIcon, CopyIcon } from "lucide-react";
+import mermaid from "mermaid";
+import { memo, type ReactNode, useEffect, useId, useState } from "react";
+import { Button } from "@/components/ui/button";
+import { copyToClipboard } from "@/lib/utils";
+
+type MermaidDiagramProps = {
+	source: string;
+	isDarkMode: boolean;
+	fallback: ReactNode;
+};
+
+let mermaidInitialized = false;
+
+function initializeMermaid() {
+	if (mermaidInitialized) return;
+
+	mermaid.initialize({
+		startOnLoad: false,
+		securityLevel: "strict",
+		htmlLabels: false,
+		flowchart: { htmlLabels: false },
+		sequence: { useMaxWidth: true },
+	});
+
+	mermaidInitialized = true;
+}
+
+function MermaidDiagramComponent({
+	source,
+	isDarkMode,
+	fallback,
+}: MermaidDiagramProps) {
+	const id = useId();
+	const [svg, setSvg] = useState<string | null>(null);
+	const [hasError, setHasError] = useState(false);
+	const [hasCopied, setHasCopied] = useState(false);
+
+	useEffect(() => {
+		let isCurrent = true;
+
+		const renderId = `mermaid-${id.replace(/[^a-zA-Z0-9_-]/g, "")}`;
+
+		setSvg(null);
+		setHasError(false);
+
+		(async () => {
+			try {
+				initializeMermaid();
+
+				// فقط theme اینجا تنظیم میشه (نه re-init کامل)
+				mermaid.initialize({
+					startOnLoad: false,
+					securityLevel: "strict",
+					htmlLabels: false,
+					theme: isDarkMode ? "dark" : "default",
+					flowchart: { htmlLabels: false },
+					sequence: { useMaxWidth: true },
+				});
+
+				await mermaid.parse(source);
+
+				const { svg } = await mermaid.render(renderId, source);
+
+				if (isCurrent) {
+					setSvg(svg);
+				}
+			} catch (error) {
+				console.error("[mermaid] Failed to render diagram", error);
+
+				if (isCurrent) {
+					setHasError(true);
+				}
+			}
+		})();
+
+		return () => {
+			isCurrent = false;
+		};
+	}, [id, isDarkMode, source]);
+
+	useEffect(() => {
+		if (!hasCopied) return;
+
+		const timer = setTimeout(() => setHasCopied(false), 2000);
+		return () => clearTimeout(timer);
+	}, [hasCopied]);
+
+	if (hasError) return fallback;
+
+	return (
+		<div className="mt-4 overflow-hidden rounded-md bg-accent">
+			<div className="flex items-center justify-between gap-4 px-4 py-2 text-sm font-semibold text-muted-foreground">
+				<span className="text-xs lowercase">mermaid</span>
+
+				<Button
+					variant="ghost"
+					size="sm"
+					className="h-8 w-8 p-0"
+					type="button"
+					onClick={async () => {
+						const ok = await copyToClipboard(source);
+						if (ok) setHasCopied(true);
+					}}
+					aria-label={hasCopied ? "Copied Mermaid source" : "Copy Mermaid source"}
+				>
+					<span className="sr-only">Copy Source</span>
+					{hasCopied ? (
+						<CheckIcon className="!size-3" />
+					) : (
+						<CopyIcon className="!size-3" />
+					)}
+				</Button>
+			</div>
+
+			<div className="bg-background/60 p-4 overflow-x-auto">
+				{svg ? (
+					// biome-ignore lint/performance/noImgElement: svg is in-memory string
+					<img
+						src={`data:image/svg+xml;charset=utf-8,${encodeURIComponent(svg)}`}
+						alt="Mermaid diagram"
+						className="mx-auto h-auto max-w-full"
+					/>
+				) : (
+					<div className="h-32 animate-pulse rounded bg-muted" />
+				)}
+			</div>
+		</div>
+	);
+}
+
+export const MermaidDiagram = memo(MermaidDiagramComponent);
\ No newline at end of file