Feat: DO-3456 extract code from markdown and use code viewer in chat component

packages/ui-components/docs/changelog.md

@@ -4,11 +4,12 @@ title: Changelog
 
 ## NEXT
 
--   Updated `Chat` component to support a non "popup" version. This can be set with the `isPopup` flag. It has also gained three other customization props `placeholder`, `chatTitle`, `loadingComponent`. 
+-   Updated `Chat` component to support a non "popup" version. This can be set with the `isPopup` flag. It has also gained three other customization props `placeholder`, `chatTitle`, `loadingComponent`.
 -   `Chat` component UI improvements according to latest design.
 -   Added two new props to `UserData` allowing to set bubble color and content in chat messages for a given user.
 -   `Textarea` component now supports `maxHeight` which when set it allows component to grow according to the size of its content.
 -   Added new `CodeViewer` component.
+-   Updated `Chat` component to use `CodeViewer` when rendering code blocks in messages. This is now also supported by default by the `Markdown` component.
 
 ## 1.10.1
 

packages/ui-components/src/chat/chat.stories.tsx

@@ -90,7 +90,28 @@ const messages = [
     },
     {
         id: 'ctDqA50c0b13FQKY0E1tf',
-        message: 'Magic!',
+        message: 'Magic! \n ```python \n def foo():\n  return \n ``` \n Some other stuff',
+        created_at: '2024-07-03T15:57:26.944Z',
+        updated_at: '2024-07-03T15:57:26.944Z',
+        user: Custom,
+    },
+    {
+        id: 'ctDqA50c0b13FQKY0E1tf',
+        message: 'Unsupported language \n ```custom \n fn foo():\n  return \n ``` \n Some other stuff',
+        created_at: '2024-07-03T15:57:26.944Z',
+        updated_at: '2024-07-03T15:57:26.944Z',
+        user: Custom,
+    },
+    {
+        id: 'ctDqA50c0b13FQKY0E1tf',
+        message: 'No language specified \n ```\nfn foo():\n  return \n ``` \n Some other stuff',
+        created_at: '2024-07-03T15:57:26.944Z',
+        updated_at: '2024-07-03T15:57:26.944Z',
+        user: Custom,
+    },
+    {
+        id: 'ctDqA50c0b13FQKY0E1tf',
+        message: 'Inline block! \n `foo = lambda x: x + 1` \n Some other stuff',
         created_at: '2024-07-03T15:57:26.944Z',
         updated_at: '2024-07-03T15:57:26.944Z',
         user: Custom,

packages/ui-components/src/code-viewer/code-viewer.tsx

@@ -114,7 +114,12 @@ function CodeViewer(props: CodeViewerProps): JSX.Element {
     }, [props.codeTheme, themeCtx.themeType]);
 
     return (
-        <CodeViewerContainer>
+        <CodeViewerContainer
+            style={{
+                ...props.style,
+            }}
+            className={props.className}
+        >
             <TopBar $isLightTheme={props.codeTheme !== 'dark'}>
                 <span>{props.language}</span>
                 {isCopied ?

packages/ui-components/src/markdown/markdown.tsx

@@ -1,7 +1,11 @@
-import ReactMarkdown, { Options } from 'react-markdown';
+import { Language } from 'prism-react-renderer';
+import * as React from 'react';
+import ReactMarkdown, { ExtraProps, Options } from 'react-markdown';
 import remarkGfm from 'remark-gfm';
 
-import styled from '@darajs/styled-components';
+import styled, { useTheme } from '@darajs/styled-components';
+
+import CodeViewer, { CodeComponentThemes } from '../code-viewer/code-viewer';
 
 interface MarkdownProps extends Options {
     /**
@@ -123,7 +127,17 @@ const CustomMarkdownWrapper = styled.div`
         font-size: 0.9rem;
     }
 
-    pre {
+    pre:first-child {
+        margin-top: 0;
+    }
+
+    pre:last-child {
+        margin-bottom: 0;
+    }
+
+    /* don't apply default styles if the pre block has the pretty code display */
+    /* stylelint-disable-next-line */
+    pre:not(.prism-code):not(:has(.prism-code)) {
         overflow-x: auto;
 
         margin-top: 1.7rem;
@@ -244,7 +258,6 @@ const CustomMarkdownWrapper = styled.div`
     img:first-child,
     figure:first-child,
     video:first-child,
-    pre:first-child,
     hr:first-child {
         margin-top: 0;
     }
@@ -262,12 +275,41 @@ const CustomMarkdownWrapper = styled.div`
     img:last-child,
     figure:last-child,
     video:last-child,
-    pre:last-child,
     hr:last-child {
         margin-bottom: 0;
     }
 `;
 
+function CodeDisplay(
+    props: React.ClassAttributes<HTMLElement> & React.HTMLAttributes<HTMLElement> & ExtraProps
+): React.ReactElement {
+    const theme = useTheme();
+    const { children, ...rest } = props;
+    const match = /language-(\w+)/.exec(props.className || '');
+
+    const parsed = React.useMemo(() => {
+        return String(children).trim().replace(/\n\n/, '\n').replace(/\n$/, '');
+    }, [children]);
+
+    // if the code block doesn't specify the language, e.g. ``` without lang or inline `code` block
+    if (!match) {
+        return <code {...rest}>{parsed}</code>;
+    }
+
+    // assume language is supported - we could check for grammar support but if unsupported
+    // prism will just render without syntax highlighting which is fine
+    return (
+        <CodeViewer
+            value={parsed}
+            language={match[1] as Language}
+            codeTheme={
+                // opposite theme
+                theme.themeType === 'dark' ? CodeComponentThemes.LIGHT : CodeComponentThemes.DARK
+            }
+        />
+    );
+}
+
 /**
  * A component for rendering markdown
  */
@@ -276,7 +318,15 @@ function Markdown(props: MarkdownProps): JSX.Element {
 
     return (
         <CustomMarkdownWrapper className={className} style={style}>
-            <ReactMarkdown {...reactMarkdownProps} remarkPlugins={reactMarkdownProps.remarkPlugins ?? [remarkGfm]}>
+            <ReactMarkdown
+                {...reactMarkdownProps}
+                components={{
+                    code: CodeDisplay,
+                    // let caller override the default components
+                    ...reactMarkdownProps.components,
+                }}
+                remarkPlugins={reactMarkdownProps.remarkPlugins ?? [remarkGfm]}
+            >
                 {markdown}
             </ReactMarkdown>
         </CustomMarkdownWrapper>