🌗 refactor: Consistent Mermaid Theming for Inline and Artifact Renderers (#12055)

* refactor: consistent theming between inline and Artifacts Mermaid Diagram * refactor: Enhance Mermaid component with improved theming and security features - Updated Mermaid component to utilize useCallback for performance optimization. - Increased maximum zoom level from 4 to 10 for better diagram visibility. - Added security level configuration to Mermaid initialization for enhanced security. - Refactored theme handling to ensure consistent theming between inline and artifact diagrams. - Introduced unit tests for Mermaid configuration to validate flowchart settings and theme behavior. * refactor: Improve theme handling in useMermaid hook - Enhanced theme variable management by merging custom theme variables with default values for dark mode. - Ensured consistent theming across Mermaid diagrams by preserving existing theme configurations while applying new defaults. * refactor: Consolidate imports in mermaid test file - Combined multiple imports from the mermaid utility into a single statement for improved readability and organization in the test file. * feat: Add subgraph title contrast adjustment for Mermaid diagrams - Introduced a utility function to enhance text visibility on subgraph titles by adjusting the fill color based on background luminance. - Updated the Mermaid component to utilize this function, ensuring better contrast in rendered SVGs. - Added comprehensive unit tests to validate the contrast adjustment logic across various scenarios. * refactor: Update MermaidHeader component for improved button accessibility and styling - Replaced Button components with TooltipAnchor for better accessibility and user experience. - Consolidated button styles into a single class for consistency. - Enhanced the layout and spacing of the header for a cleaner appearance. * fix: hex color handling and improve contrast adjustment in Mermaid diagrams - Updated hexLuminance function to support 3-character hex shorthand by expanding it to 6 characters. - Refined the fixSubgraphTitleContrast function to avoid double semicolons in style attributes and ensure proper fill color adjustments based on background luminance. - Added unit tests to validate the handling of 3-character hex fills and the prevention of double semicolons in text styles. * chore: Simplify Virtual Scrolling Performance tests by removing performance timing checks - Removed performance timing checks and associated console logs from tests handling 1000 and 5000 agents. - Focused tests on verifying the correct rendering of virtual list items without measuring render time.
2026-03-26 17:46:34 +01:00 · 2026-03-04 08:25:57 -05:00 · 2026-03-04 08:25:57 -05:00 · f1eabdbdb7
commit f1eabdbdb7
parent 6ebee069c7
8 changed files with 671 additions and 309 deletions
--- a/client/src/utils/tests/mermaid.test.ts
+++ b/client/src/utils/tests/mermaid.test.ts
@ -0,0 +1,172 @@
+import {
+  fixSubgraphTitleContrast,
+  artifactFlowchartConfig,
+  inlineFlowchartConfig,
+  getMermaidFiles,
+} from '~/utils/mermaid';
+
+const makeSvg = (clusters: string): Element => {
+  const parser = new DOMParser();
+  const doc = parser.parseFromString(
+    `<svg xmlns="http://www.w3.org/2000/svg">${clusters}</svg>`,
+    'image/svg+xml',
+  );
+  return doc.querySelector('svg')!;
+};
+
+describe('mermaid config', () => {
+  describe('flowchart config invariants', () => {
+    it('inlineFlowchartConfig must have htmlLabels: false for blob URL <img> rendering', () => {
+      expect(inlineFlowchartConfig.htmlLabels).toBe(false);
+    });
+
+    it('artifactFlowchartConfig must have htmlLabels: true for direct DOM injection', () => {
+      expect(artifactFlowchartConfig.htmlLabels).toBe(true);
+    });
+
+    it('both configs share the same base layout settings', () => {
+      expect(inlineFlowchartConfig.curve).toBe(artifactFlowchartConfig.curve);
+      expect(inlineFlowchartConfig.nodeSpacing).toBe(artifactFlowchartConfig.nodeSpacing);
+      expect(inlineFlowchartConfig.rankSpacing).toBe(artifactFlowchartConfig.rankSpacing);
+      expect(inlineFlowchartConfig.padding).toBe(artifactFlowchartConfig.padding);
+    });
+  });
+
+  describe('getMermaidFiles', () => {
+    const content = 'graph TD\n  A-->B';
+
+    it('produces dark theme files when isDarkMode is true', () => {
+      const files = getMermaidFiles(content, true);
+      expect(files['/components/ui/MermaidDiagram.tsx']).toContain('theme: "dark"');
+      expect(files['mermaid.css']).toContain('#212121');
+    });
+
+    it('produces neutral theme files when isDarkMode is false', () => {
+      const files = getMermaidFiles(content, false);
+      expect(files['/components/ui/MermaidDiagram.tsx']).toContain('theme: "neutral"');
+      expect(files['mermaid.css']).toContain('#FFFFFF');
+    });
+
+    it('defaults to dark mode when isDarkMode is omitted', () => {
+      const files = getMermaidFiles(content);
+      expect(files['/components/ui/MermaidDiagram.tsx']).toContain('theme: "dark"');
+    });
+
+    it('includes securityLevel in generated component', () => {
+      const files = getMermaidFiles(content, true);
+      expect(files['/components/ui/MermaidDiagram.tsx']).toContain('securityLevel: "strict"');
+    });
+
+    it('includes all required file keys', () => {
+      const files = getMermaidFiles(content, true);
+      expect(files['diagram.mmd']).toBe(content);
+      expect(files['App.tsx']).toBeDefined();
+      expect(files['index.tsx']).toBeDefined();
+      expect(files['/components/ui/MermaidDiagram.tsx']).toBeDefined();
+      expect(files['mermaid.css']).toBeDefined();
+    });
+
+    it('uses artifact flowchart config with htmlLabels: true', () => {
+      const files = getMermaidFiles(content, true);
+      expect(files['/components/ui/MermaidDiagram.tsx']).toContain('"htmlLabels": true');
+    });
+
+    it('does not inject custom themeVariables into generated component', () => {
+      const darkFiles = getMermaidFiles(content, true);
+      const lightFiles = getMermaidFiles(content, false);
+      expect(darkFiles['/components/ui/MermaidDiagram.tsx']).not.toContain('themeVariables');
+      expect(lightFiles['/components/ui/MermaidDiagram.tsx']).not.toContain('themeVariables');
+    });
+
+    it('handles empty content', () => {
+      const files = getMermaidFiles('', true);
+      expect(files['diagram.mmd']).toBe('# No mermaid diagram content provided');
+    });
+  });
+
+  describe('fixSubgraphTitleContrast', () => {
+    it('darkens title text on light subgraph backgrounds (fill attribute)', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#FFF9C4"/><g class="cluster-label"><text fill="#E0E0E0">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toContain('fill: #1a1a1a');
+    });
+
+    it('darkens title text on light subgraph backgrounds (inline style fill)', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect style="fill: #FFF9C4; stroke: #F9A825"/><g class="cluster-label"><text>Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toContain('fill: #1a1a1a');
+    });
+
+    it('lightens title text on dark subgraph backgrounds', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#1f2020"/><g class="cluster-label"><text fill="#222222">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toContain('fill: #f0f0f0');
+    });
+
+    it('leaves title text alone when contrast is already good', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#FFF9C4"/><g class="cluster-label"><text fill="#333333">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toBeNull();
+    });
+
+    it('skips clusters without a rect', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><g class="cluster-label"><text fill="#E0E0E0">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toBeNull();
+    });
+
+    it('skips clusters with non-hex fills', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="rgb(255,249,196)"/><g class="cluster-label"><text fill="#E0E0E0">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toBeNull();
+    });
+
+    it('sets dark fill when text has no explicit fill on light backgrounds', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect style="fill:#FFF9C4"/><g class="cluster-label"><text>Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toContain('fill: #1a1a1a');
+    });
+
+    it('preserves existing text style when appending fill override', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#FFF9C4"/><g class="cluster-label"><text style="font-size: 14px" fill="#E0E0E0">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      const style = svg.querySelector('text')!.getAttribute('style')!;
+      expect(style).toContain('font-size: 14px');
+      expect(style).toContain('fill: #1a1a1a');
+    });
+
+    it('handles 3-char hex shorthand fills', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#FFC"/><g class="cluster-label"><text fill="#EEE">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      expect(svg.querySelector('text')!.getAttribute('style')).toContain('fill: #1a1a1a');
+    });
+
+    it('avoids double semicolons when existing style has trailing semicolon', () => {
+      const svg = makeSvg(
+        '<g class="cluster"><rect fill="#FFF9C4"/><g class="cluster-label"><text style="font-size: 14px;" fill="#E0E0E0">Title</text></g></g>',
+      );
+      fixSubgraphTitleContrast(svg);
+      const style = svg.querySelector('text')!.getAttribute('style')!;
+      expect(style).not.toContain(';;');
+      expect(style).toContain('fill: #1a1a1a');
+    });
+  });
+});
--- a/client/src/utils/mermaid.ts
+++ b/client/src/utils/mermaid.ts
@ -1,15 +1,157 @@
 import dedent from 'dedent';

-const mermaid = dedent(`import React, { useEffect, useRef, useState } from "react";
+interface MermaidButtonStyles {
+  bg: string;
+  bgHover: string;
+  border: string;
+  text: string;
+  textSecondary: string;
+  shadow: string;
+  divider: string;
+}
+
+const darkButtonStyles: MermaidButtonStyles = {
+  bg: 'rgba(40, 40, 40, 0.95)',
+  bgHover: 'rgba(60, 60, 60, 0.95)',
+  border: '1px solid rgba(255, 255, 255, 0.1)',
+  text: '#D1D5DB',
+  textSecondary: '#9CA3AF',
+  shadow: '0 2px 8px rgba(0, 0, 0, 0.4)',
+  divider: 'rgba(255, 255, 255, 0.1)',
+};
+
+const lightButtonStyles: MermaidButtonStyles = {
+  bg: 'rgba(255, 255, 255, 0.95)',
+  bgHover: 'rgba(243, 244, 246, 0.95)',
+  border: '1px solid rgba(0, 0, 0, 0.1)',
+  text: '#374151',
+  textSecondary: '#6B7280',
+  shadow: '0 2px 8px rgba(0, 0, 0, 0.1)',
+  divider: 'rgba(0, 0, 0, 0.1)',
+};
+
+const getButtonStyles = (isDarkMode: boolean): MermaidButtonStyles =>
+  isDarkMode ? darkButtonStyles : lightButtonStyles;
+
+const baseFlowchartConfig = {
+  curve: 'basis' as const,
+  nodeSpacing: 50,
+  rankSpacing: 50,
+  diagramPadding: 8,
+  useMaxWidth: true,
+  padding: 15,
+  wrappingWidth: 200,
+};
+
+/** Artifact renderer injects SVG directly into the DOM where foreignObject works */
+const artifactFlowchartConfig = {
+  ...baseFlowchartConfig,
+  htmlLabels: true,
+};
+
+/** Inline renderer converts SVG to a blob URL <img>; browsers block foreignObject in that context */
+const inlineFlowchartConfig = {
+  ...baseFlowchartConfig,
+  htmlLabels: false,
+};
+
+export { inlineFlowchartConfig, artifactFlowchartConfig };
+
+/** Perceived luminance (0 = black, 1 = white) via BT.601 luma coefficients */
+const hexLuminance = (hex: string): number => {
+  let h = hex.replace('#', '');
+  if (h.length === 3) {
+    h = h[0] + h[0] + h[1] + h[1] + h[2] + h[2];
+  }
+  if (h.length < 6) {
+    return -1;
+  }
+  const r = parseInt(h.slice(0, 2), 16) / 255;
+  const g = parseInt(h.slice(2, 4), 16) / 255;
+  const b = parseInt(h.slice(4, 6), 16) / 255;
+  return 0.299 * r + 0.587 * g + 0.114 * b;
+};
+
+/**
+ * Fixes subgraph title text contrast in mermaid SVGs rendered with htmlLabels: false.
+ * When a subgraph has an explicit light fill via `style` directives, the title <text>
+ * gets its fill from a CSS rule (.cluster-label text / .cluster text) set to titleColor.
+ * In dark mode, titleColor is light, producing invisible text on light backgrounds.
+ * This walks cluster groups and overrides the text fill attribute when contrast is poor.
+ */
+export const fixSubgraphTitleContrast = (svgElement: Element): void => {
+  const clusters = svgElement.querySelectorAll('g.cluster');
+  for (const cluster of clusters) {
+    const rect = cluster.querySelector(':scope > rect, :scope > polygon');
+    if (!rect) {
+      continue;
+    }
+
+    const inlineStyle = rect.getAttribute('style') || '';
+    const attrFill = rect.getAttribute('fill') || '';
+    const styleFillMatch = inlineStyle.match(/fill\s*:\s*(#[0-9a-fA-F]{3,8})/);
+    const hex = styleFillMatch?.[1] ?? (attrFill.startsWith('#') ? attrFill : '');
+    if (!hex) {
+      continue;
+    }
+
+    const bgLum = hexLuminance(hex);
+    if (bgLum < 0) {
+      continue;
+    }
+
+    const textElements = cluster.querySelectorAll(
+      ':scope > g.cluster-label text, :scope > text, :scope > g > text',
+    );
+    for (const textEl of textElements) {
+      const textFill = textEl.getAttribute('fill') || '';
+      const textStyle = textEl.getAttribute('style') || '';
+      const textStyleFill = textStyle.match(/fill\s*:\s*(#[0-9a-fA-F]{3,8})/);
+      const currentHex = textStyleFill?.[1] ?? (textFill.startsWith('#') ? textFill : '');
+      const isLightBg = bgLum > 0.5;
+
+      let newFill = '';
+      if (!currentHex) {
+        if (isLightBg) {
+          newFill = '#1a1a1a';
+        }
+      } else {
+        const textLum = hexLuminance(currentHex);
+        if (textLum < 0) {
+          continue;
+        }
+        if (isLightBg && textLum > 0.5) {
+          newFill = '#1a1a1a';
+        } else if (!isLightBg && textLum < 0.4) {
+          newFill = '#f0f0f0';
+        }
+      }
+
+      if (newFill) {
+        let sep = '';
+        if (textStyle) {
+          sep = textStyle.trimEnd().endsWith(';') ? ' ' : '; ';
+        }
+        textEl.setAttribute('style', `${textStyle}${sep}fill: ${newFill}`);
+      }
+    }
+  }
+};
+
+const buildMermaidComponent = (
+  mermaidTheme: string,
+  bgColor: string,
+  btnStyles: MermaidButtonStyles,
+) =>
+  dedent(`import React, { useEffect, useRef, useState, useCallback } from "react";
 import {
  TransformWrapper,
  TransformComponent,
  ReactZoomPanPinchRef,
 } from "react-zoom-pan-pinch";
 import mermaid from "mermaid";
-import { Button } from "/components/ui/button";

-const ZoomIn = () => (
+const ZoomInIcon = () => (
  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
    <circle cx="11" cy="11" r="8"/>
    <line x1="21" x2="16.65" y1="21" y2="16.65"/>
@ -18,7 +160,7 @@ const ZoomIn = () => (
  </svg>
 );

-const ZoomOut = () => (
+const ZoomOutIcon = () => (
  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
    <circle cx="11" cy="11" r="8"/>
    <line x1="21" x2="16.65" y1="21" y2="16.65"/>
@ -26,90 +168,147 @@ const ZoomOut = () => (
  </svg>
 );

-const RefreshCw = () => (
+const ResetIcon = () => (
  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
-    <path d="M3 12a9 9 0 0 1 9-9 9.75 9.75 0 0 1 6.74 2.74L21 8"/>
-    <path d="M21 3v5h-5"/>
-    <path d="M21 12a9 9 0 0 1-9 9 9.75 9.75 0 0 1-6.74-2.74L3 16"/>
-    <path d="M8 16H3v5"/>
+    <path d="M1 4v6h6"/>
+    <path d="M3.51 15a9 9 0 1 0 2.13-9.36L1 10"/>
  </svg>
 );

+const CopyIcon = () => (
+  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
+    <rect width="14" height="14" x="8" y="8" rx="2" ry="2"/>
+    <path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/>
+  </svg>
+);
+
+const CheckIcon = () => (
+  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
+    <polyline points="20 6 9 17 4 12"/>
+  </svg>
+);
+
+const btnBase = {
+  display: "inline-flex",
+  alignItems: "center",
+  justifyContent: "center",
+  width: "28px",
+  height: "28px",
+  borderRadius: "4px",
+  background: "transparent",
+  border: "none",
+  color: "${btnStyles.text}",
+  cursor: "pointer",
+  padding: "6px",
+  transition: "background 0.15s ease",
+};
+
+const btnHover = {
+  ...btnBase,
+  background: "${btnStyles.bgHover}",
+};
+
+const toolbarStyle = {
+  position: "absolute",
+  bottom: "12px",
+  right: "12px",
+  display: "flex",
+  alignItems: "center",
+  gap: "2px",
+  padding: "4px",
+  borderRadius: "8px",
+  background: "${btnStyles.bg}",
+  boxShadow: "${btnStyles.shadow}",
+  backdropFilter: "blur(8px)",
+  border: "${btnStyles.border}",
+  zIndex: 10,
+};
+
+const dividerStyle = {
+  width: "1px",
+  height: "16px",
+  background: "${btnStyles.divider}",
+  margin: "0 4px",
+};
+
+const zoomTextStyle = {
+  minWidth: "3rem",
+  textAlign: "center",
+  fontSize: "12px",
+  color: "${btnStyles.textSecondary}",
+  userSelect: "none",
+  fontFamily: "system-ui, -apple-system, sans-serif",
+};
+
 interface MermaidDiagramProps {
  content: string;
 }

+interface IconButtonProps {
+  onClick: () => void;
+  children: React.ReactNode;
+  title: string;
+  disabled?: boolean;
+}
+
+const IconButton = ({ onClick, children, title, disabled = false }: IconButtonProps) => {
+  const [hovered, setHovered] = useState(false);
+  return (
+    <button
+      type="button"
+      onClick={onClick}
+      title={title}
+      aria-label={title}
+      disabled={disabled}
+      style={{
+        ...(hovered && !disabled ? btnHover : btnBase),
+        opacity: disabled ? 0.4 : 1,
+        cursor: disabled ? "default" : "pointer",
+      }}
+      onMouseEnter={() => setHovered(true)}
+      onMouseLeave={() => setHovered(false)}
+    >
+      {children}
+    </button>
+  );
+};
+
+const ZOOM_STEP = 0.1;
+const MIN_SCALE = 0.1;
+const MAX_SCALE = 10;
+
 const MermaidDiagram: React.FC<MermaidDiagramProps> = ({ content }) => {
  const mermaidRef = useRef<HTMLDivElement>(null);
  const transformRef = useRef<ReactZoomPanPinchRef>(null);
  const [isRendered, setIsRendered] = useState(false);
+  const [zoomLevel, setZoomLevel] = useState(100);
+  const [copied, setCopied] = useState(false);

  useEffect(() => {
    mermaid.initialize({
      startOnLoad: false,
-      theme: "base",
-      themeVariables: {
-        background: "#282C34",
-        primaryColor: "#333842",
-        secondaryColor: "#333842",
-        tertiaryColor: "#333842",
-        primaryTextColor: "#ABB2BF",
-        secondaryTextColor: "#ABB2BF",
-        lineColor: "#636D83",
-        fontSize: "16px",
-        nodeBorder: "#636D83",
-        mainBkg: '#282C34',
-        altBackground: '#282C34',
-        textColor: '#ABB2BF',
-        edgeLabelBackground: '#282C34',
-        clusterBkg: '#282C34',
-        clusterBorder: "#636D83",
-        labelBoxBkgColor: "#333842",
-        labelBoxBorderColor: "#636D83",
-        labelTextColor: "#ABB2BF",
-      },
-      flowchart: {
-        curve: "basis",
-        nodeSpacing: 50,
-        rankSpacing: 50,
-        diagramPadding: 8,
-        htmlLabels: true,
-        useMaxWidth: true,
-        padding: 15,
-        wrappingWidth: 200,
-      },
+      theme: "${mermaidTheme}",
+      securityLevel: "strict",
+      flowchart: ${JSON.stringify(artifactFlowchartConfig, null, 8)},
    });

    const renderDiagram = async () => {
-      if (mermaidRef.current) {
-        try {
-          const { svg } = await mermaid.render("mermaid-diagram", content);
-          mermaidRef.current.innerHTML = svg;
+      if (!mermaidRef.current) {
+        return;
+      }
+      try {
+        const { svg } = await mermaid.render("mermaid-diagram", content);
+        mermaidRef.current.innerHTML = svg;

-          const svgElement = mermaidRef.current.querySelector("svg");
-          if (svgElement) {
-            svgElement.style.width = "100%";
-            svgElement.style.height = "100%";
-
-            const pathElements = svgElement.querySelectorAll("path");
-            pathElements.forEach((path) => {
-              path.style.strokeWidth = "1.5px";
-            });
-
-            const rectElements = svgElement.querySelectorAll("rect");
-            rectElements.forEach((rect) => {
-              const parent = rect.parentElement;
-              if (parent && parent.classList.contains("node")) {
-                rect.style.stroke = "#636D83";
-                rect.style.strokeWidth = "1px";
-              } else {
-                rect.style.stroke = "none";
-              }
-            });
-          }
-          setIsRendered(true);
-        } catch (error) {
-          console.error("Mermaid rendering error:", error);
+        const svgElement = mermaidRef.current.querySelector("svg");
+        if (svgElement) {
+          svgElement.style.width = "100%";
+          svgElement.style.height = "100%";
+        }
+        setIsRendered(true);
+      } catch (error) {
+        console.error("Mermaid rendering error:", error);
+        if (mermaidRef.current) {
          mermaidRef.current.innerHTML = "Error rendering diagram";
        }
      }
@ -118,72 +317,90 @@ const MermaidDiagram: React.FC<MermaidDiagramProps> = ({ content }) => {
    renderDiagram();
  }, [content]);

-  const centerAndFitDiagram = () => {
+  const centerAndFitDiagram = useCallback(() => {
    if (transformRef.current && mermaidRef.current) {
      const { centerView, zoomToElement } = transformRef.current;
-      zoomToElement(mermaidRef.current as HTMLElement);
+      zoomToElement(mermaidRef.current);
      centerView(1, 0);
+      setZoomLevel(100);
    }
-  };
+  }, []);

  useEffect(() => {
    if (isRendered) {
      centerAndFitDiagram();
    }
-  }, [isRendered]);
+  }, [isRendered, centerAndFitDiagram]);

-  const handlePanning = () => {
-    if (transformRef.current) {
-      const { state, instance } = transformRef.current;
-      if (!state) {
-        return;
-      }
-      const { scale, positionX, positionY } = state;
-      const { wrapperComponent, contentComponent } = instance;
-
-      if (wrapperComponent && contentComponent) {
-        const wrapperRect = wrapperComponent.getBoundingClientRect();
-        const contentRect = contentComponent.getBoundingClientRect();
-        const maxX = wrapperRect.width - contentRect.width * scale;
-        const maxY = wrapperRect.height - contentRect.height * scale;
-
-        let newX = positionX;
-        let newY = positionY;
-
-        if (newX > 0) {
-          newX = 0;
-        }
-        if (newY > 0) {
-          newY = 0;
-        }
-        if (newX < maxX) {
-          newX = maxX;
-        }
-        if (newY < maxY) {
-          newY = maxY;
-        }
-
-        if (newX !== positionX || newY !== positionY) {
-          instance.setTransformState(scale, newX, newY);
-        }
-      }
+  const handleTransform = useCallback((ref) => {
+    if (ref && ref.state) {
+      setZoomLevel(Math.round(ref.state.scale * 100));
    }
-  };
+  }, []);
+
+  const handleCopy = useCallback(() => {
+    navigator.clipboard.writeText(content).then(() => {
+      setCopied(true);
+      setTimeout(() => setCopied(false), 2000);
+    }).catch(() => {});
+  }, [content]);
+
+  const handlePanning = useCallback(() => {
+    if (!transformRef.current) {
+      return;
+    }
+    const { state, instance } = transformRef.current;
+    if (!state) {
+      return;
+    }
+    const { scale, positionX, positionY } = state;
+    const { wrapperComponent, contentComponent } = instance;
+
+    if (!wrapperComponent || !contentComponent) {
+      return;
+    }
+
+    const wrapperRect = wrapperComponent.getBoundingClientRect();
+    const contentRect = contentComponent.getBoundingClientRect();
+    const maxX = wrapperRect.width - contentRect.width * scale;
+    const maxY = wrapperRect.height - contentRect.height * scale;
+
+    let newX = positionX;
+    let newY = positionY;
+
+    if (newX > 0) {
+      newX = 0;
+    }
+    if (newY > 0) {
+      newY = 0;
+    }
+    if (newX < maxX) {
+      newX = maxX;
+    }
+    if (newY < maxY) {
+      newY = maxY;
+    }
+
+    if (newX !== positionX || newY !== positionY) {
+      instance.setTransformState(scale, newX, newY);
+    }
+  }, []);

  return (
-    <div className="relative h-screen w-screen cursor-move bg-[#282C34] p-5">
+    <div style={{ position: "relative", height: "100vh", width: "100vw", cursor: "move", padding: "20px", backgroundColor: "${bgColor}" }}>
      <TransformWrapper
        ref={transformRef}
        initialScale={1}
-        minScale={0.1}
-        maxScale={10}
+        minScale={MIN_SCALE}
+        maxScale={MAX_SCALE}
        limitToBounds={false}
        centerOnInit={true}
        initialPositionY={0}
-        wheel={{ step: 0.1 }}
+        wheel={{ step: ZOOM_STEP }}
        panning={{ velocityDisabled: true }}
        alignmentAnimation={{ disabled: true }}
        onPanning={handlePanning}
+        onTransformed={handleTransform}
      >
        {({ zoomIn, zoomOut }) => (
          <>
@ -204,24 +421,22 @@ const MermaidDiagram: React.FC<MermaidDiagramProps> = ({ content }) => {
                }}
              />
            </TransformComponent>
-            <div className="absolute bottom-2 right-2 flex space-x-2">
-              <Button onClick={() => zoomIn(0.1)} variant="outline" size="icon">
-                <ZoomIn />
-              </Button>
-              <Button
-                onClick={() => zoomOut(0.1)}
-                variant="outline"
-                size="icon"
-              >
-                <ZoomOut />
-              </Button>
-              <Button
-                onClick={centerAndFitDiagram}
-                variant="outline"
-                size="icon"
-              >
-                <RefreshCw />
-              </Button>
+            <div style={toolbarStyle}>
+              <IconButton onClick={() => zoomOut(ZOOM_STEP)} title="Zoom out">
+                <ZoomOutIcon />
+              </IconButton>
+              <span style={zoomTextStyle}>{zoomLevel}%</span>
+              <IconButton onClick={() => zoomIn(ZOOM_STEP)} title="Zoom in">
+                <ZoomInIcon />
+              </IconButton>
+              <div style={dividerStyle} />
+              <IconButton onClick={centerAndFitDiagram} title="Reset view">
+                <ResetIcon />
+              </IconButton>
+              <div style={dividerStyle} />
+              <IconButton onClick={handleCopy} title="Copy code">
+                {copied ? <CheckIcon /> : <CopyIcon />}
+              </IconButton>
            </div>
          </>
        )}
@ -242,13 +457,16 @@ export default App = () => (
 `);
 };

-const mermaidCSS = `
+export const getMermaidFiles = (content: string, isDarkMode = true) => {
+  const mermaidTheme = isDarkMode ? 'dark' : 'neutral';
+  const btnStyles = getButtonStyles(isDarkMode);
+  const bgColor = isDarkMode ? '#212121' : '#FFFFFF';
+  const mermaidCSS = `
 body {
-  background-color: #282C34;
+  background-color: ${bgColor};
 }
 `;

-export const getMermaidFiles = (content: string) => {
  return {
    'diagram.mmd': content || '# No mermaid diagram content provided',
    'App.tsx': wrapMermaidDiagram(content),
@ -262,7 +480,7 @@ import App from "./App";
 const root = createRoot(document.getElementById("root"));
 root.render(<App />);
 ;`),
-    '/components/ui/MermaidDiagram.tsx': mermaid,
+    '/components/ui/MermaidDiagram.tsx': buildMermaidComponent(mermaidTheme, bgColor, btnStyles),
    'mermaid.css': mermaidCSS,
  };
 };