fix(ui): stabilize mermaid diagram rendering while streaming (#1676)

Signed-off-by: Petr Kadlec <[email protected]>

Author: kapetr

apps/agentstack-ui/src/components/MarkdownContent/MarkdownContent.tsx

@@ -14,14 +14,15 @@ import type { PluggableList } from 'unified';
 import { components, type ExtendedComponents } from './components';
 import { Code } from './components/Code';
 import { MermaidDiagram } from './components/MermaidDiagram';
+import { MermaidProvider } from './contexts';
 import classes from './MarkdownContent.module.scss';
 import { rehypePlugins } from './rehype';
 import { remarkPlugins } from './remark';
 import { urlTransform } from './utils';
 
 export interface MarkdownContentProps {
   codeBlocksExpanded?: boolean;
-  showMermaidDiagrams?: boolean;
+  isStreaming?: boolean;
   children?: string;
   className?: string;
   remarkPlugins?: PluggableList;
@@ -30,7 +31,7 @@ export interface MarkdownContentProps {
 
 export function MarkdownContent({
   codeBlocksExpanded,
-  showMermaidDiagrams,
+  isStreaming,
   className,
   remarkPlugins: remarkPluginsProps,
   components: componentsProps,
@@ -40,24 +41,26 @@ export function MarkdownContent({
     () => ({
       ...components,
       code: ({ ...props }) => <Code {...props} forceExpand={codeBlocksExpanded} />,
-      mermaidDiagram: ({ ...props }) => <MermaidDiagram {...props} showDiagram={showMermaidDiagrams} />,
+      mermaidDiagram: (props) => <MermaidDiagram {...props} isStreaming={isStreaming} />,
       ...componentsProps,
     }),
-    [codeBlocksExpanded, showMermaidDiagrams, componentsProps],
+    [codeBlocksExpanded, componentsProps, isStreaming],
   );
 
   const extendedRemarkPlugins = useMemo(() => [...remarkPlugins, ...(remarkPluginsProps ?? [])], [remarkPluginsProps]);
 
   return (
-    <div className={clsx(classes.root, className)}>
-      <Markdown
-        rehypePlugins={rehypePlugins}
-        remarkPlugins={extendedRemarkPlugins}
-        components={extendedComponents}
-        urlTransform={urlTransform}
-      >
-        {children}
-      </Markdown>
-    </div>
+    <MermaidProvider>
+      <div className={clsx(classes.root, className)}>
+        <Markdown
+          rehypePlugins={rehypePlugins}
+          remarkPlugins={extendedRemarkPlugins}
+          components={extendedComponents}
+          urlTransform={urlTransform}
+        >
+          {children}
+        </Markdown>
+      </div>
+    </MermaidProvider>
   );
 }

apps/agentstack-ui/src/components/MarkdownContent/components/MermaidDiagram.module.scss

@@ -12,3 +12,20 @@
 .diagram {
   text-align: center;
 }
+
+.loading {
+  display: flex;
+  justify-content: center;
+  align-items: center;
+  min-block-size: rem(120px);
+
+  :global(.cds--inline-loading) {
+    justify-content: center;
+  }
+}
+
+.error {
+  :global(.cds--inline-notification__subtitle) {
+    @include line-clamp(2);
+  }
+}

apps/agentstack-ui/src/components/MarkdownContent/components/MermaidDiagram.tsx

@@ -3,48 +3,59 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { InlineLoading, InlineNotification } from '@carbon/react';
 import mermaid from 'mermaid';
-import { type HTMLAttributes, useEffect, useId, useState } from 'react';
+import { type HTMLAttributes, useEffect, useId } from 'react';
 import type { ExtraProps } from 'react-markdown';
 
 import { useTheme } from '#contexts/Theme/index.ts';
 import { Theme } from '#contexts/Theme/types.ts';
 
+import { useMermaid } from '../contexts';
 import { Code } from './Code';
 import classes from './MermaidDiagram.module.scss';
 
 export type MermaidDiagramProps = HTMLAttributes<HTMLElement> &
-  ExtraProps & {
-    showDiagram?: boolean;
-  };
+  ExtraProps & { mermaidIndex?: number; isStreaming?: boolean };
 
-export function MermaidDiagram({ showDiagram = true, children }: MermaidDiagramProps) {
+export function MermaidDiagram({ children, mermaidIndex, isStreaming }: MermaidDiagramProps) {
   const id = useId();
-  const [diagram, setDiagram] = useState<string | null>(null);
-
   const { theme } = useTheme();
+  const { diagrams, setDiagram } = useMermaid();
+
+  if (mermaidIndex === undefined) {
+    console.error('MermaidDiagram component requires a `mermaidIndex` prop.');
+  }
+  const index = mermaidIndex ?? 0;
+
+  const diagram = diagrams.get(index);
 
   useEffect(() => {
-    mermaid.initialize({ startOnLoad: false, theme: theme === Theme.Dark ? 'dark' : 'default' });
+    mermaid.initialize({
+      startOnLoad: false,
+      theme: theme === Theme.Dark ? 'dark' : 'default',
+      suppressErrorRendering: true,
+    });
   }, [theme]);
 
   useEffect(() => {
     let isMounted = true;
 
     async function renderDiagram() {
-      if (!showDiagram || typeof children !== 'string') {
+      if (typeof children !== 'string') {
         return;
       }
 
       try {
         const { svg } = await mermaid.render(id, children);
 
         if (isMounted) {
-          setDiagram(svg);
+          setDiagram(index, svg);
         }
       } catch (error) {
-        if (isMounted) {
+        if (isMounted && !isStreaming) {
           console.warn(error);
+          setDiagram(index, error instanceof Error ? error : new Error('Unknown error rendering Mermaid diagram'));
         }
       }
     }
@@ -54,13 +65,27 @@ export function MermaidDiagram({ showDiagram = true, children }: MermaidDiagramP
     return () => {
       isMounted = false;
     };
-  }, [showDiagram, children, theme, id]);
+  }, [children, theme, id, setDiagram, index, isStreaming]);
 
   return (
     <div className={classes.root}>
       <Code className="language-mermaid">{children}</Code>
 
-      {showDiagram && diagram && <div dangerouslySetInnerHTML={{ __html: diagram }} className={classes.diagram} />}
+      {typeof diagram === 'string' ? (
+        <div dangerouslySetInnerHTML={{ __html: diagram }} className={classes.diagram} />
+      ) : diagram instanceof Error && !isStreaming ? (
+        <InlineNotification
+          kind="error"
+          title={'Failed to render Mermaid diagram'}
+          subtitle={diagram.message}
+          lowContrast
+          className={classes.error}
+        />
+      ) : (
+        <div className={classes.loading}>
+          <InlineLoading description="Rendering diagram..." />
+        </div>
+      )}
     </div>
   );
 }

apps/agentstack-ui/src/components/MarkdownContent/contexts/MermaidProvider.tsx

@@ -0,0 +1,32 @@
+/**
+ * Copyright 2025 © BeeAI a Series of LF Projects, LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { PropsWithChildren } from 'react';
+import { useCallback, useMemo, useState } from 'react';
+
+import type { MermaidContextValue } from './mermaid-context';
+import { MermaidContext } from './mermaid-context';
+
+export function MermaidProvider({ children }: PropsWithChildren) {
+  const [diagrams, setDiagrams] = useState<Map<number, string | Error>>(new Map());
+
+  const setDiagram = useCallback((index: number, svg: string) => {
+    setDiagrams((prev) => {
+      const newMap = new Map(prev);
+      newMap.set(index, svg);
+      return newMap;
+    });
+  }, []);
+
+  const value = useMemo<MermaidContextValue>(
+    () => ({
+      diagrams,
+      setDiagram,
+    }),
+    [diagrams, setDiagram],
+  );
+
+  return <MermaidContext.Provider value={value}>{children}</MermaidContext.Provider>;
+}

apps/agentstack-ui/src/components/MarkdownContent/contexts/index.ts

@@ -0,0 +1,7 @@
+/**
+ * Copyright 2025 © BeeAI a Series of LF Projects, LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export { MermaidProvider } from './MermaidProvider';
+export { useMermaid } from './useMermaid';

apps/agentstack-ui/src/components/MarkdownContent/contexts/mermaid-context.ts

@@ -0,0 +1,13 @@
+/**
+ * Copyright 2025 © BeeAI a Series of LF Projects, LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+'use client';
+import { createContext } from 'react';
+
+export const MermaidContext = createContext<MermaidContextValue | null>(null);
+
+export interface MermaidContextValue {
+  diagrams: Map<number, string | Error>;
+  setDiagram: (index: number, value: string | Error) => void;
+}

apps/agentstack-ui/src/components/MarkdownContent/contexts/useMermaid.ts

@@ -0,0 +1,19 @@
+/**
+ * Copyright 2025 © BeeAI a Series of LF Projects, LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+'use client';
+
+import { useContext } from 'react';
+
+import { MermaidContext } from './mermaid-context';
+
+export function useMermaid() {
+  const context = useContext(MermaidContext);
+
+  if (!context) {
+    throw new Error('useMermaid must be used within MermaidProvider');
+  }
+
+  return context;
+}

apps/agentstack-ui/src/components/MarkdownContent/remark/remarkMermaid.ts

@@ -9,12 +9,19 @@ import { visit } from 'unist-util-visit';
 
 export function remarkMermaid() {
   return (tree: Root) => {
+    let mermaidIndex = 0;
+
     visit(tree, 'code', (node: Code) => {
       if (node.lang === 'mermaid') {
         node.data = {
           ...node.data,
           hName: 'mermaidDiagram',
+          hProperties: {
+            mermaidIndex,
+          },
         };
+
+        mermaidIndex++;
       }
     });
   };

apps/agentstack-ui/src/modules/messages/components/MessageContent.tsx

@@ -9,7 +9,6 @@ import { memo } from 'react';
 import type { UIMessage } from '#modules/messages/types.ts';
 import { ChatMarkdownContent } from '#modules/runs/components/ChatMarkdownContent/ChatMarkdownContent.tsx';
 
-import { useAgentRun } from '../../runs/contexts/agent-run';
 import { Role } from '../api/types';
 import { checkMessageStatus, getMessageContent, getMessageSecret, getMessageSources, isAgentMessage } from '../utils';
 import classes from './MessageContent.module.scss';
@@ -20,8 +19,6 @@ interface Props {
 }
 
 export const MessageContent = memo(function MessageContent({ message }: Props) {
-  const { isPending } = useAgentRun();
-
   const content = getMessageContent(message);
   const form = message.role === Role.User ? message.form : null;
   const auth = message.role === Role.User ? message.auth : null;
@@ -45,8 +42,8 @@ export const MessageContent = memo(function MessageContent({ message }: Props) {
       <ChatMarkdownContent
         className={classes.root}
         sources={sources}
-        codeBlocksExpanded={isPending}
-        showMermaidDiagrams={!isPending}
+        codeBlocksExpanded={status?.isInProgress}
+        isStreaming={status?.isInProgress}
       >
         {content}
       </ChatMarkdownContent>

apps/agentstack-ui/src/modules/runs/components/RunOutputBox.tsx

@@ -31,7 +31,7 @@ export function RunOutputBox({ isPending, text, downloadFileName, sources, child
       )}
 
       {text && (
-        <ChatMarkdownContent sources={sources} codeBlocksExpanded={isPending} showMermaidDiagrams={!isPending}>
+        <ChatMarkdownContent sources={sources} codeBlocksExpanded={isPending} isStreaming={isPending}>
           {text}
         </ChatMarkdownContent>
       )}