feat(markdown): fix Mermaid integration with error handling and copy button

studio/frontend/package.json

@@ -66,7 +66,7 @@
     "remark-gfm": "^4.0.1",
     "shadcn": "^3.8.4",
     "sonner": "^2.0.7",
-    "streamdown": "^2.2.0",
+    "streamdown": "^2.3.0",
     "tailwind-merge": "^3.4.0",
     "tailwindcss": "^4.1.18",
     "tw-animate-css": "^1.4.0",
@@ -90,4 +90,4 @@
     "typescript-eslint": "^8.55.0",
     "vite": "^7.3.1"
   }
-}
+}

studio/frontend/src/components/assistant-ui/markdown-text.tsx

@@ -1,14 +1,83 @@
 "use client";
 
 import { INTERNAL, useMessagePartText } from "@assistant-ui/react";
+import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { Copy02Icon, Tick02Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
 import { code } from "@streamdown/code";
 import { math } from "@streamdown/math";
 import { mermaid } from "@streamdown/mermaid";
-import { Streamdown } from "streamdown";
+import { Block, type BlockProps, Streamdown } from "streamdown";
+import { useEffect, useRef, useState } from "react";
 import "katex/dist/katex.min.css";
 
 const { withSmoothContextProvider, useSmoothStatus } = INTERNAL;
 
+function getMermaidSource(blockContent: string): string | null {
+  const source = blockContent.match(/```mermaid\s*([\s\S]*?)```/i)?.[1]?.trim();
+  return source && source.length > 0 ? source : null;
+}
+
+const COPY_RESET_MS = 2000;
+
+function MermaidCopyButton({ source }: { source: string }) {
+  const [copied, setCopied] = useState(false);
+  const resetTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    return () => {
+      if (resetTimeoutRef.current) {
+        clearTimeout(resetTimeoutRef.current);
+      }
+    };
+  }, []);
+
+  return (
+    <button
+      type="button"
+      className="absolute top-3.5 right-20 z-20 cursor-pointer text-muted-foreground transition-all hover:text-foreground"
+      title="Copy Mermaid source"
+      onClick={() => {
+        if (!copyToClipboard(source)) return;
+        setCopied(true);
+        if (resetTimeoutRef.current) {
+          clearTimeout(resetTimeoutRef.current);
+        }
+        resetTimeoutRef.current = setTimeout(() => {
+          setCopied(false);
+          resetTimeoutRef.current = null;
+        }, COPY_RESET_MS);
+      }}
+    >
+      <HugeiconsIcon icon={copied ? Tick02Icon : Copy02Icon} className="size-5" />
+    </button>
+  );
+}
+
+function StreamdownBlock(props: BlockProps) {
+  const hasMermaidFence = props.content.includes("```mermaid");
+  const mermaidSource = getMermaidSource(props.content);
+
+  if (props.isIncomplete && hasMermaidFence) {
+    return (
+      <div className="my-4 flex h-48 items-center justify-center rounded-xl border border-border bg-muted/30 text-sm text-muted-foreground animate-pulse">
+        Loading diagram...
+      </div>
+    );
+  }
+
+  if (mermaidSource) {
+    return (
+      <div className="relative">
+        <Block {...props} />
+        <MermaidCopyButton source={mermaidSource} />
+      </div>
+    );
+  }
+
+  return <Block {...props} />;
+}
+
 const MarkdownTextImpl = () => {
   const { text } = useMessagePartText();
   const status = useSmoothStatus();
@@ -19,8 +88,16 @@ const MarkdownTextImpl = () => {
         mode="streaming"
         isAnimating={status.type === "running"}
         plugins={{ code, math, mermaid }}
-        controls={true}
+        controls={{
+          mermaid: {
+            fullscreen: true,
+            download: true,
+            copy: false,
+            panZoom: true,
+          },
+        }}
         shikiTheme={["github-light", "github-dark"]}
+        BlockComponent={StreamdownBlock}
       >
         {text}
       </Streamdown>

studio/frontend/src/components/markdown/mermaid-error.tsx

@@ -0,0 +1,28 @@
+import type { MermaidErrorComponentProps } from "streamdown";
+
+function hasSlashComment(chart: string): boolean {
+  return /(^|[^:])\/\/.*/m.test(chart);
+}
+
+export function MermaidError({
+  error,
+  chart,
+  retry,
+}: MermaidErrorComponentProps) {
+  return (
+    <div className="my-4 rounded-lg border border-red-300 bg-red-50 p-3 text-red-800">
+      <p className="text-sm font-semibold">Mermaid render failed</p>
+      <p className="mt-1 break-words font-mono text-xs">{error}</p>
+      {hasSlashComment(chart) ? (
+        <p className="mt-1 text-xs">Hint: Mermaid comments use `%%`, not `//`.</p>
+      ) : null}
+      <button
+        type="button"
+        onClick={retry}
+        className="mt-2 rounded border border-red-300 px-2 py-1 text-xs hover:bg-red-100"
+      >
+        Retry
+      </button>
+    </div>
+  );
+}