fix(ui): hide mermaid diagram during streaming (#1669)

Signed-off-by: Petr Bulánek <[email protected]>

Author: PetrBulanek

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

@@ -13,13 +13,15 @@ import type { PluggableList } from 'unified';
 
 import { components, type ExtendedComponents } from './components';
 import { Code } from './components/Code';
+import { MermaidDiagram } from './components/MermaidDiagram';
 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;
   children?: string;
   className?: string;
   remarkPlugins?: PluggableList;
@@ -28,6 +30,7 @@ export interface MarkdownContentProps {
 
 export function MarkdownContent({
   codeBlocksExpanded,
+  showMermaidDiagrams,
   className,
   remarkPlugins: remarkPluginsProps,
   components: componentsProps,
@@ -37,9 +40,10 @@ export function MarkdownContent({
     () => ({
       ...components,
       code: ({ ...props }) => <Code {...props} forceExpand={codeBlocksExpanded} />,
+      mermaidDiagram: ({ ...props }) => <MermaidDiagram {...props} showDiagram={showMermaidDiagrams} />,
       ...componentsProps,
     }),
-    [codeBlocksExpanded, componentsProps],
+    [codeBlocksExpanded, showMermaidDiagrams, componentsProps],
   );
 
   const extendedRemarkPlugins = useMemo(() => [...remarkPlugins, ...(remarkPluginsProps ?? [])], [remarkPluginsProps]);

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

@@ -13,47 +13,54 @@ import { Theme } from '#contexts/Theme/types.ts';
 import { Code } from './Code';
 import classes from './MermaidDiagram.module.scss';
 
-export type MermaidDiagramProps = HTMLAttributes<HTMLElement> & ExtraProps;
+export type MermaidDiagramProps = HTMLAttributes<HTMLElement> &
+  ExtraProps & {
+    showDiagram?: boolean;
+  };
 
-export function MermaidDiagram({ children }: MermaidDiagramProps) {
+export function MermaidDiagram({ showDiagram = true, children }: MermaidDiagramProps) {
   const id = useId();
   const [diagram, setDiagram] = useState<string | null>(null);
 
   const { theme } = useTheme();
 
   useEffect(() => {
-    let cancelled = false;
+    mermaid.initialize({ startOnLoad: false, theme: theme === Theme.Dark ? 'dark' : 'default' });
+  }, [theme]);
 
-    (async () => {
-      if (typeof children !== 'string') {
+  useEffect(() => {
+    let isMounted = true;
+
+    async function renderDiagram() {
+      if (!showDiagram || typeof children !== 'string') {
         return;
       }
 
       try {
-        mermaid.initialize({ startOnLoad: false, theme: theme === Theme.Dark ? 'dark' : 'default' });
-
         const { svg } = await mermaid.render(id, children);
 
-        if (!cancelled) {
+        if (isMounted) {
           setDiagram(svg);
         }
       } catch (error) {
-        if (!cancelled) {
+        if (isMounted) {
           console.warn(error);
         }
       }
-    })();
+    }
+
+    renderDiagram();
 
     return () => {
-      cancelled = true;
+      isMounted = false;
     };
-  }, [children, theme, id]);
+  }, [showDiagram, children, theme, id]);
 
   return (
     <div className={classes.root}>
       <Code className="language-mermaid">{children}</Code>
 
-      {diagram && <div dangerouslySetInnerHTML={{ __html: diagram }} className={classes.diagram} />}
+      {showDiagram && diagram && <div dangerouslySetInnerHTML={{ __html: diagram }} className={classes.diagram} />}
     </div>
   );
 }

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

@@ -42,7 +42,12 @@ export const MessageContent = memo(function MessageContent({ message }: Props) {
     }
 
     return (
-      <ChatMarkdownContent className={classes.root} sources={sources} codeBlocksExpanded={isPending}>
+      <ChatMarkdownContent
+        className={classes.root}
+        sources={sources}
+        codeBlocksExpanded={isPending}
+        showMermaidDiagrams={!isPending}
+      >
         {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}>
+        <ChatMarkdownContent sources={sources} codeBlocksExpanded={isPending} showMermaidDiagrams={!isPending}>
           {text}
         </ChatMarkdownContent>
       )}