Switch Assistant to Streamdown (#39281)

* switch to streamdown

* Merge master resolve conflicts

---------

Co-authored-by: Joshen Lim <joshenlimek@gmail.com>

apps/studio/components/ui/AIAssistantPanel/Message.Display.tsx

@@ -5,7 +5,8 @@ import { ProfileImage as ProfileImageDisplay } from 'components/ui/ProfileImage'
 import { useProfileNameAndPicture } from 'lib/profile'
 import { cn } from 'ui'
 import { useMessageInfoContext } from './Message.Context'
-import { MessageMarkdown, MessagePartSwitcher } from './Message.Parts'
+import { MessageMarkdown } from './MessageMarkdown'
+import { MessagePartSwitcher } from './Message.Parts'
 
 function MessageDisplayProfileImage() {
   const { username, avatarUrl } = useProfileNameAndPicture()
@@ -75,7 +76,7 @@ function MessageDisplayTextMessage({
       id={id}
       isLoading={isLoading}
       readOnly={readOnly}
-      className="prose prose-sm max-w-none break-words"
+      className="prose prose-sm max-w-none break-words prose-h2:font-medium"
     >
       {children}
     </MessageMarkdown>

apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx

@@ -1,12 +1,8 @@
 import { UIMessage as VercelMessage } from '@ai-sdk/react'
 import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai'
 import { BrainIcon, CheckIcon, Loader2 } from 'lucide-react'
-import { useMemo, type PropsWithChildren } from 'react'
-import ReactMarkdown from 'react-markdown'
-import { type Components } from 'react-markdown/lib/ast-to-react'
-import remarkGfm from 'remark-gfm'
 
-import { cn, markdownComponents } from 'ui'
+import { cn } from 'ui'
 import { DisplayBlockRenderer } from './DisplayBlockRenderer'
 import { EdgeFunctionRenderer } from './EdgeFunctionRenderer'
 import { Tool } from './elements/Tool'
@@ -16,71 +12,7 @@ import {
   deployEdgeFunctionOutputSchema,
   parseExecuteSqlChartResult,
 } from './Message.utils'
-import {
-  Heading3,
-  Hyperlink,
-  InlineCode,
-  ListItem,
-  MarkdownPre,
-  OrderedList,
-} from './MessageMarkdown'
-
-const baseMarkdownComponents: Partial<Components> = {
-  ol: OrderedList,
-  li: ListItem,
-  h3: Heading3,
-  code: InlineCode,
-  a: Hyperlink,
-  img: ({ src }) => <span className="text-foreground-light font-mono">[Image: {src}]</span>,
-}
-
-export function MessageMarkdown({
-  id,
-  isLoading,
-  readOnly,
-  className,
-  children,
-}: PropsWithChildren<{
-  id: string
-  isLoading: boolean
-  readOnly?: boolean
-  className?: string
-}>) {
-  const markdownSource = useMemo(() => {
-    if (typeof children === 'string') {
-      return children
-    }
-
-    if (Array.isArray(children)) {
-      return children.filter((child): child is string => typeof child === 'string').join('')
-    }
-
-    return ''
-  }, [children])
-
-  const allMarkdownComponents: Partial<Components> = useMemo(
-    () => ({
-      ...markdownComponents,
-      ...baseMarkdownComponents,
-      pre: ({ children }) => (
-        <MarkdownPre id={id} isLoading={isLoading} readOnly={readOnly}>
-          {children}
-        </MarkdownPre>
-      ),
-    }),
-    [id, isLoading, readOnly]
-  )
-
-  return (
-    <ReactMarkdown
-      className={className}
-      remarkPlugins={[remarkGfm]}
-      components={allMarkdownComponents}
-    >
-      {markdownSource}
-    </ReactMarkdown>
-  )
-}
+import { MessageMarkdown } from './MessageMarkdown'
 
 function MessagePartText({ textPart }: { textPart: TextUIPart }) {
   const { id, isLoading, readOnly, isUserMessage, state } = useMessageInfoContext()
@@ -91,7 +23,7 @@ function MessagePartText({ textPart }: { textPart: TextUIPart }) {
       isLoading={isLoading}
       readOnly={readOnly}
       className={cn(
-        'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
+        'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h2:font-medium prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
         isUserMessage && 'text-foreground [&>p]:font-medium',
         state === 'editing' && 'animate-pulse'
       )}

apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx

@@ -1,6 +1,15 @@
-import { Loader2 } from 'lucide-react'
+import dynamic from 'next/dynamic'
 import Link from 'next/link'
-import { memo, ReactNode, useEffect, useMemo, useRef } from 'react'
+import {
+  isValidElement,
+  memo,
+  ReactNode,
+  useEffect,
+  useMemo,
+  useRef,
+  type ReactElement,
+} from 'react'
+import type { StreamdownProps } from 'streamdown'
 
 import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig'
 import {
@@ -8,6 +17,7 @@ import {
   cn,
   CodeBlock,
   CodeBlockLang,
+  markdownComponents,
   Dialog,
   DialogClose,
   DialogContent,
@@ -23,29 +33,34 @@ import { CollapsibleCodeBlock } from './CollapsibleCodeBlock'
 import { DisplayBlockRenderer } from './DisplayBlockRenderer'
 import { defaultUrlTransform } from './Message.utils'
 
-export const OrderedList = memo(({ children }: { children: ReactNode }) => (
+const Streamdown = dynamic<StreamdownProps>(
+  () => import('streamdown').then((mod) => mod.Streamdown),
+  { ssr: false }
+)
+
+export const OrderedList = memo(({ children }: { children?: ReactNode }) => (
   <ol className="flex flex-col gap-y-4">{children}</ol>
 ))
 OrderedList.displayName = 'OrderedList'
 
-export const ListItem = memo(({ children }: { children: ReactNode }) => (
+export const ListItem = memo(({ children }: { children?: ReactNode }) => (
   <li className="[&>pre]:mt-2">{children}</li>
 ))
 ListItem.displayName = 'ListItem'
 
-export const Heading3 = memo(({ children }: { children: ReactNode }) => (
+export const Heading3 = memo(({ children }: { children?: ReactNode }) => (
   <h3 className="underline">{children}</h3>
 ))
 Heading3.displayName = 'Heading3'
 
 export const InlineCode = memo(
-  ({ className, children }: { className?: string; children: ReactNode }) => (
+  ({ className, children }: { className?: string; children?: ReactNode }) => (
     <code className={cn('text-xs', className)}>{children}</code>
   )
 )
 InlineCode.displayName = 'InlineCode'
 
-export const Hyperlink = memo(({ href, children }: { href?: string; children: ReactNode }) => {
+export const Hyperlink = memo(({ href, children }: { href?: string; children?: ReactNode }) => {
   const isExternalURL = !href?.startsWith('https://supabase.com/dashboard')
   const safeUrl = defaultUrlTransform(href ?? '')
   const isSafeUrl = safeUrl.length > 0
@@ -103,6 +118,62 @@ export const Hyperlink = memo(({ href, children }: { href?: string; children: Re
 })
 Hyperlink.displayName = 'Hyperlink'
 
+const baseMarkdownComponents = {
+  ol: OrderedList,
+  li: ListItem,
+  h3: Heading3,
+  code: InlineCode,
+  a: Hyperlink,
+  img: ({ src }: JSX.IntrinsicElements['img']) => (
+    <span className="text-foreground-light font-mono">[Image: {src}]</span>
+  ),
+}
+
+export function MessageMarkdown({
+  id,
+  isLoading,
+  readOnly,
+  className,
+  children,
+}: {
+  id: string
+  isLoading: boolean
+  readOnly?: boolean
+  className?: string
+  children: ReactNode
+}) {
+  const markdownSource = useMemo(() => {
+    if (typeof children === 'string') {
+      return children
+    }
+
+    if (Array.isArray(children)) {
+      return children.filter((child): child is string => typeof child === 'string').join('')
+    }
+
+    return ''
+  }, [children])
+
+  const allMarkdownComponents = useMemo(
+    () => ({
+      ...markdownComponents,
+      ...baseMarkdownComponents,
+      pre: (props: JSX.IntrinsicElements['pre']) => (
+        <MarkdownPre id={id} isLoading={isLoading} readOnly={readOnly}>
+          {props.children}
+        </MarkdownPre>
+      ),
+    }),
+    [id, isLoading, readOnly]
+  )
+
+  return (
+    <Streamdown className={className} components={allMarkdownComponents}>
+      {markdownSource}
+    </Streamdown>
+  )
+}
+
 export const MarkdownPre = ({
   children,
   id,
@@ -123,8 +194,16 @@ export const MarkdownPre = ({
     cumulative: false,
   })
 
-  const language = children[0].props.className?.replace('language-', '') || 'sql'
-  const rawContent = children[0].props.children[0]
+  const childArray = Array.isArray(children) ? children : [children]
+  const codeElement = childArray.find((child): child is ReactElement => isValidElement(child))
+  const codeProps = codeElement?.props || {}
+  const language = codeProps.className?.replace('language-', '') || 'sql'
+  const codeChildren = codeProps.children
+  const rawContent = Array.isArray(codeChildren)
+    ? codeChildren.map((node) => (typeof node === 'string' ? node : '')).join('')
+    : typeof codeChildren === 'string'
+      ? codeChildren
+      : ''
   const propsMatch = rawContent.match(/(?:--|\/\/)\s*props:\s*(\{[^}]+\})/)
 
   const snippetProps: AssistantSnippetProps = useMemo(() => {
@@ -155,6 +234,10 @@ export const MarkdownPre = ({
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [snippetProps])
 
+  if (!codeElement) {
+    return <pre className="w-auto overflow-x-auto not-prose my-4">{children}</pre>
+  }
+
   return (
     <div className="w-auto overflow-x-hidden not-prose my-4 ">
       {language === 'edge' ? (
@@ -167,11 +250,6 @@ export const MarkdownPre = ({
       ) : language === 'sql' ? (
         readOnly ? (
           <CollapsibleCodeBlock value={cleanContent} language="sql" hideLineNumbers />
-        ) : isLoading ? (
-          <div className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2">
-            <Loader2 className="w-4 h-4 animate-spin" />
-            Writing SQL...
-          </div>
         ) : (
           <DisplayBlockRenderer
             messageId={id}

apps/studio/package.json

@@ -100,7 +100,7 @@
     "monaco-editor": "0.52.2",
     "next": "catalog:",
     "next-themes": "^0.3.0",
-    "nuqs": "^2.7.0",
+    "nuqs": "2.7.1",
     "openai": "^4.75.1",
     "openapi-fetch": "0.12.4",
     "papaparse": "^5.3.1",
@@ -136,6 +136,7 @@
     "sonner": "^1.5.0",
     "sql-formatter": "^15.0.0",
     "sqlstring": "^2.3.2",
+    "streamdown": "^1.3.0",
     "tus-js-client": "^4.1.0",
     "ui": "workspace:*",
     "ui-patterns": "workspace:*",