fixup! review comments, render mermaid errors

Signed-off-by: Petr Kadlec <petr@puradesign.cz>

Author: kapetr

.vscode/settings.json

@@ -0,0 +1,16 @@
+{
+  "editor.codeActionsOnSave": {
+    "source.fixAll.eslint": "always"
+  },
+  "[markdown]": {
+    "editor.formatOnSave": false
+  },
+  "[yaml]": {
+    "editor.formatOnSave": false
+  },
+  "search.exclude": {
+    "apps/agentstack-ui/dist": true,
+    "apps/agentstack-ui/.next": true
+  },
+  "stylelint.validate": ["css", "scss"]
+}

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

@@ -22,6 +22,7 @@ import { urlTransform } from './utils';
 
 export interface MarkdownContentProps {
   codeBlocksExpanded?: boolean;
+  isStreaming?: boolean;
   children?: string;
   className?: string;
   remarkPlugins?: PluggableList;
@@ -30,6 +31,7 @@ export interface MarkdownContentProps {
 
 export function MarkdownContent({
   codeBlocksExpanded,
+  isStreaming,
   className,
   remarkPlugins: remarkPluginsProps,
   components: componentsProps,
@@ -39,10 +41,10 @@ export function MarkdownContent({
     () => ({
       ...components,
       code: ({ ...props }) => <Code {...props} forceExpand={codeBlocksExpanded} />,
-      mermaidDiagram: (props) => <MermaidDiagram {...props} />,
+      mermaidDiagram: (props) => <MermaidDiagram {...props} isStreaming={isStreaming} />,
       ...componentsProps,
     }),
-    [codeBlocksExpanded, componentsProps],
+    [codeBlocksExpanded, componentsProps, isStreaming],
   );
 
   const extendedRemarkPlugins = useMemo(() => [...remarkPlugins, ...(remarkPluginsProps ?? [])], [remarkPluginsProps]);

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

@@ -23,3 +23,9 @@
     justify-content: center;
   }
 }
+
+.error {
+  :global(.cds--inline-notification__subtitle) {
+    @include line-clamp(2);
+  }
+}

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

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { InlineLoading } from '@carbon/react';
+import { InlineLoading, InlineNotification } from '@carbon/react';
 import mermaid from 'mermaid';
 import { type HTMLAttributes, useEffect, useId } from 'react';
 import type { ExtraProps } from 'react-markdown';
@@ -15,9 +15,10 @@ import { useMermaid } from '../contexts';
 import { Code } from './Code';
 import classes from './MermaidDiagram.module.scss';
 
-export type MermaidDiagramProps = HTMLAttributes<HTMLElement> & ExtraProps & { mermaidIndex?: number };
+export type MermaidDiagramProps = HTMLAttributes<HTMLElement> &
+  ExtraProps & { mermaidIndex?: number; isStreaming?: boolean };
 
-export function MermaidDiagram({ children, mermaidIndex }: MermaidDiagramProps) {
+export function MermaidDiagram({ children, mermaidIndex, isStreaming }: MermaidDiagramProps) {
   const id = useId();
   const { theme } = useTheme();
   const { diagrams, setDiagram } = useMermaid();
@@ -29,8 +30,14 @@ export function MermaidDiagram({ children, mermaidIndex }: MermaidDiagramProps)
 
   const diagram = diagrams.get(index);
 
+  console.log({ index, isStreaming, diagram });
+
   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(() => {
@@ -48,8 +55,9 @@ export function MermaidDiagram({ children, mermaidIndex }: MermaidDiagramProps)
           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'));
         }
       }
     }
@@ -59,14 +67,22 @@ export function MermaidDiagram({ children, mermaidIndex }: MermaidDiagramProps)
     return () => {
       isMounted = false;
     };
-  }, [children, theme, id, setDiagram, index]);
+  }, [children, theme, id, setDiagram, index, isStreaming]);
 
   return (
     <div className={classes.root}>
       <Code className="language-mermaid">{children}</Code>
 
-      {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..." />

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

@@ -10,7 +10,7 @@ import type { MermaidContextValue } from './mermaid-context';
 import { MermaidContext } from './mermaid-context';
 
 export function MermaidProvider({ children }: PropsWithChildren) {
-  const [diagrams, setDiagrams] = useState<Map<number, string>>(new Map());
+  const [diagrams, setDiagrams] = useState<Map<number, string | Error>>(new Map());
 
   const setDiagram = useCallback((index: number, svg: string) => {
     setDiagrams((prev) => {

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

@@ -8,6 +8,6 @@ import { createContext } from 'react';
 export const MermaidContext = createContext<MermaidContextValue | null>(null);
 
 export interface MermaidContextValue {
-  diagrams: Map<number, string>;
-  setDiagram: (index: number, svg: string) => void;
+  diagrams: Map<number, string | Error>;
+  setDiagram: (index: number, value: string | Error) => void;
 }

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

@@ -9,12 +9,10 @@ import type { PluggableList } from 'unified';
 
 import { remarkExternalLink } from './remarkExternalLink';
 import { remarkMermaid } from './remarkMermaid';
-import { remarkMermaidIndex } from './remarkMermaidIndex';
 
 export const remarkPlugins = [
   remarkGfm,
   [remarkMath, { singleDollarTextMath: false }],
-  remarkMermaidIndex,
   remarkMermaid,
   remarkExternalLink,
 ] satisfies PluggableList;

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/components/MarkdownContent/remark/remarkMermaidIndex.ts

@@ -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;
@@ -42,7 +39,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={status?.isInProgress}
+        isStreaming={status?.isInProgress}
+      >
         {content}
       </ChatMarkdownContent>
     );