feat(chat): add mermaid chart support to markdown display #103

src/components/MarkdownDisplay.tsx

@@ -1,6 +1,6 @@
 import 'katex/dist/katex.min.css';
 import { all as languages } from 'lowlight';
-import React, { memo, useMemo } from 'react';
+import React, { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { LuCopy, LuPlay } from 'react-icons/lu';
 import Markdown, { ExtraProps } from 'react-markdown';
@@ -12,8 +12,9 @@ import remarkMath from 'remark-math';
 import { useAppContext } from '../context/app';
 import { useChatContext } from '../context/chat';
 import { CanvasType } from '../types';
-import { classNames, copyStr } from '../utils';
+import { copyStr } from '../utils';
 import { IntlIconButton } from './common';
+import MermaidChart from './MermaidChart';
 
 export default memo(function MarkdownDisplay({ content }: { content: string }) {
   const preprocessedContent = useMemo(
@@ -55,98 +56,114 @@ const CustomPre: React.ElementType<
     ExtraProps & { origContent: string }
 > = ({ className, children, node, origContent }) => {
   const { t } = useTranslation();
-  const {
-    config: { pyIntepreterEnabled },
-  } = useAppContext();
-  const { setCanvasData } = useChatContext();
 
-  const showActionButtons = useMemo(() => {
+  const { language, code } = useMemo(() => {
     const startOffset = node?.position?.start.offset;
     const endOffset = node?.position?.end.offset;
-    if (!startOffset || !endOffset) return false;
-    return true;
-  }, [node?.position]);
+    if (!startOffset || !endOffset) {
+      return {
+        language: '',
+        code: '',
+      };
+    }
+    return {
+      language:
+        origContent
+          .substring(startOffset, endOffset)
+          .match(/^```([^\n]+)\n/)?.[1] ?? '',
+      code: getCodeContent(origContent.substring(startOffset, endOffset)),
+    };
+  }, [node?.position, origContent]);
 
-  const codeLanguage = useMemo(() => {
-    const startOffset = node?.position?.start.offset;
-    const endOffset = node?.position?.end.offset;
-    if (!startOffset || !endOffset) return '';
+  if (!!language && !!code && language.toLowerCase() === 'mermaid') {
+    return <MermaidChart className="mermaid-diagram" code={code} />;
+  }
 
-    return (
-      origContent
-        .substring(startOffset, endOffset)
-        .match(/^```([^\n]+)\n/)?.[1] ?? ''
-    );
-  }, [node?.position, origContent]);
+  return (
+    <div className="hljs" aria-label={t('chatScreen.ariaLabels.codeBlock')}>
+      <CodeBlockToolbar language={language} code={code} />
+
+      <pre className={className} {...node?.properties}>
+        <CodeLanguageDisplay language={language} />
+        {children}
+      </pre>
+    </div>
+  );
+};
+
+function CodeBlockToolbar({
+  language,
+  code,
+}: {
+  language: string;
+  code: string;
+}) {
+  const { t } = useTranslation();
+  const {
+    config: { pyIntepreterEnabled },
+  } = useAppContext();
+  const { setCanvasData } = useChatContext();
 
   const canRunCode = useMemo(
-    () => pyIntepreterEnabled && codeLanguage.toLowerCase() === 'python',
-    [pyIntepreterEnabled, codeLanguage]
+    () => pyIntepreterEnabled && language.toLowerCase() === 'python',
+    [pyIntepreterEnabled, language]
   );
 
-  const handleCopy = () => {
-    const startOffset = node?.position?.start.offset;
-    const endOffset = node?.position?.end.offset;
-    if (!startOffset || !endOffset) return;
+  const handleCopy = useCallback(() => {
+    if (!code) return;
+    copyStr(code);
+  }, [code]);
 
-    copyStr(getCodeContent(origContent.substring(startOffset, endOffset)));
-  };
-  const handleRun = () => {
-    const startOffset = node?.position?.start.offset;
-    const endOffset = node?.position?.end.offset;
-    if (!startOffset || !endOffset) return;
+  const handleRun = useCallback(() => {
+    if (!code) return;
+
+    if (language.toLowerCase() === 'python') {
+      setCanvasData({
+        type: CanvasType.PY_INTERPRETER,
+        content: code,
+      });
+    }
+  }, [code, language, setCanvasData]);
 
-    setCanvasData({
-      type: CanvasType.PY_INTERPRETER,
-      content: getCodeContent(origContent.substring(startOffset, endOffset)),
-    });
-  };
+  if (!language || !code) return null;
 
   return (
-    <div className="hljs" aria-label={t('chatScreen.ariaLabels.codeBlock')}>
-      {showActionButtons && (
-        <div
-          className={classNames({
-            'hljs sticky h-0 z-[1] text-right p-0': true,
-            'display-none': !node?.position,
-          })}
-        >
-          {canRunCode && (
-            <IntlIconButton
-              className="btn btn-ghost w-8 h-8 p-0"
-              t={t}
-              titleKey="chatScreen.titles.run"
-              ariaLabelKey="chatScreen.ariaLabels.runCode"
-              icon={LuPlay}
-              onClick={handleRun}
-            />
-          )}
-          <IntlIconButton
-            className="btn btn-ghost w-8 h-8 p-0"
-            t={t}
-            titleKey="chatScreen.titles.copy"
-            ariaLabelKey="chatScreen.ariaLabels.copyContent"
-            icon={LuCopy}
-            onClick={handleCopy}
-          />
-        </div>
+    <div className="hljs sticky h-0 z-[1] text-right p-0">
+      {canRunCode && (
+        <IntlIconButton
+          className="btn btn-ghost w-8 h-8 p-0"
+          t={t}
+          titleKey="chatScreen.titles.run"
+          ariaLabelKey="chatScreen.ariaLabels.runCode"
+          icon={LuPlay}
+          onClick={handleRun}
+        />
       )}
+      <IntlIconButton
+        className="btn btn-ghost w-8 h-8 p-0"
+        t={t}
+        titleKey="chatScreen.titles.copy"
+        ariaLabelKey="chatScreen.ariaLabels.copyContent"
+        icon={LuCopy}
+        onClick={handleCopy}
+      />
+    </div>
+  );
+}
 
-      <pre className={className} {...node?.properties}>
-        {codeLanguage && (
-          <div
-            className="text-sm ml-2"
-            aria-label={t('chatScreen.ariaLabels.codeLanguage')}
-          >
-            {codeLanguage}
-          </div>
-        )}
+function CodeLanguageDisplay({ language }: { language: string }) {
+  const { t } = useTranslation();
+  if (!language) return null;
 
-        {children}
-      </pre>
+  return (
+    <div
+      className="text-sm ml-2"
+      aria-label={t('chatScreen.ariaLabels.codeLanguage')}
+    >
+      {language}
     </div>
   );
-};
+}
 
 /**
  * The part below is copied and adapted from:

src/components/MermaidChart.tsx

@@ -0,0 +1,88 @@
+import { useEffect, useMemo, useRef, useState } from 'react';
+import { isDev } from '../config';
+
+declare global {
+  interface Window {
+    mermaid: {
+      initialize: (config: Record<string, unknown>) => void;
+      render: (
+        id: string,
+        text: string,
+        svgContainingElement?: Element
+      ) => Promise<{
+        svg: string;
+        diagramType: string;
+        bindFunctions?: (element: Element) => void;
+      }>;
+    };
+  }
+}
+
+type MermaidTheme = 'default' | 'forest' | 'dark' | 'neutral';
+
+interface MermaidDrawerProps {
+  className?: string;
+  code: string;
+  theme?: MermaidTheme;
+}
+
+export default function MermaidChart({
+  className,
+  code,
+  theme = 'default',
+}: MermaidDrawerProps) {
+  const chartRef = useRef<HTMLDivElement | null>(null);
+  const [mermaidLoaded, setMermaidLoaded] = useState(false);
+
+  const renderId = useMemo(
+    () =>
+      `mermaid-diagram-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
+    []
+  );
+
+  useEffect(() => {
+    if (window.mermaid) {
+      window.mermaid.initialize({ startOnLoad: false, theme });
+      setMermaidLoaded(true);
+      return;
+    }
+    const script = document.createElement('script');
+    script.src = 'https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.min.js';
+    script.onload = () => {
+      if (!window.mermaid) return;
+      window.mermaid.initialize({ startOnLoad: false, theme });
+      setMermaidLoaded(true);
+      if (isDev) console.debug('Mermaid is loaded');
+    };
+    document.body.appendChild(script);
+    return () => {
+      if (script.parentNode) script.parentNode.removeChild(script);
+    };
+  }, [theme]);
+
+  useEffect(() => {
+    if (!mermaidLoaded || !chartRef.current || !code) return;
+    if (isDev) console.debug('Mermaid rendering:', renderId);
+    window.mermaid
+      .render(renderId, code)
+      .then(({ svg, diagramType }) => {
+        if (chartRef.current) {
+          chartRef.current.innerHTML = svg;
+          chartRef.current.ariaLabel = `Mermaid ${diagramType} chart`;
+        }
+      })
+      .catch((error) => {
+        console.error('Mermaid rendering error:', error);
+        if (!chartRef.current) return;
+        chartRef.current.innerHTML = `<pre>${code}</pre>`;
+        chartRef.current.ariaLabel =
+          'Mermaid diagram (render failed, showing raw code)';
+      });
+  }, [code, mermaidLoaded, renderId]);
+
+  return (
+    <div ref={chartRef} className={className} role="img" aria-label="Diagram">
+      <pre className="mermaid">${code}</pre>
+    </div>
+  );
+}