improvement(chat): reduce streaming render work and completion delay

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/chat-content.tsx

@@ -11,7 +11,7 @@ import {
   useState,
 } from 'react'
 import type { Nodes } from 'hast'
-import { defaultRehypePlugins, Streamdown } from 'streamdown'
+import { defaultRehypePlugins, defaultRemarkPlugins, Streamdown } from 'streamdown'
 import 'streamdown/styles.css'
 // prismjs core must load before its language components — they register on the
 // global `Prism` it installs (on `window`/`global`); fixes SSR + client order.
@@ -35,6 +35,8 @@ import {
   isInlineFileReference,
 } from '@/lib/mothership/chat/inline-image-reference'
 import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
+import { HighlightedLines } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines'
+import { remarkPlainText } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text'
 import {
   SourceChip,
   sourceLabel,
@@ -66,6 +68,11 @@ const LANG_ALIASES: Record<string, string> = {
   py: 'python',
 }
 
+const MARKDOWN_REMARK_PLUGINS = [
+  ...Object.values(defaultRemarkPlugins),
+  remarkPlainText,
+] as const satisfies NonNullable<ComponentPropsWithoutRef<typeof Streamdown>['remarkPlugins']>
+
 const PROSE_CLASSES = cn(
   'prose prose-base dark:prose-invert max-w-none',
   'font-[family-name:var(--font-inter)] antialiased break-words tracking-[0]',
@@ -308,10 +315,9 @@ const MARKDOWN_COMPONENTS = {
           />
         </div>
         <div className='code-editor-theme bg-[var(--surface-5)] dark:bg-[var(--code-bg)]'>
-          <pre
-            className='m-0 overflow-x-auto whitespace-pre p-4 font-mono text-[var(--text-primary)] text-small leading-[21px]'
-            dangerouslySetInnerHTML={{ __html: html }}
-          />
+          <pre className='m-0 overflow-x-auto whitespace-pre p-4 font-mono text-[var(--text-primary)] text-small leading-[21px]'>
+            <HighlightedLines html={html} />
+          </pre>
         </div>
       </div>
     )
@@ -704,6 +710,7 @@ function ChatContentInner({
                     isAnimating={streamingTree}
                     components={MARKDOWN_COMPONENTS}
                     rehypePlugins={imageRehypePlugins}
+                    remarkPlugins={MARKDOWN_REMARK_PLUGINS}
                   >
                     {group.markdown}
                   </Streamdown>

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines.test.tsx

@@ -0,0 +1,83 @@
+/** @vitest-environment jsdom */
+import { act } from 'react'
+import Prism from 'prismjs'
+import { createRoot } from 'react-dom/client'
+import 'prismjs/components/prism-typescript'
+import { describe, expect, it } from 'vitest'
+import {
+  HighlightedLines,
+  splitHighlightedLines,
+} from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines'
+
+/** Compare per-character token ancestry, ignoring unstyled line wrappers. */
+function styledCharacters(element: Element): Array<[string, string]> {
+  const result: Array<[string, string]> = []
+  function visit(node: Node, classes: string[]) {
+    if (node.nodeType === Node.TEXT_NODE) {
+      for (const char of node.textContent ?? '') {
+        if (char !== '\n') result.push([char, classes.join('/')])
+      }
+    } else if (node instanceof Element) {
+      const nextClasses = node.className ? [...classes, node.className] : classes
+      for (const child of node.childNodes) visit(child, nextClasses)
+    }
+  }
+  visit(element, [])
+  return result
+}
+
+describe('highlighted code lines', () => {
+  it('preserves full-fence Prism text and token styles across multiline constructs', () => {
+    const cases = [
+      `const value = \`first\nsecond \${1 + 2}\nthird\`\n`,
+      '/* a comment\nwith <script>alert(1)</script>\n*/\nconst x = "safe"',
+      '<div title="a &amp; b">\n  <!-- multi\n  line -->\n</div>',
+      'const x = "<img src=x onerror=alert(1)>"\n\nconst y = 1',
+    ] as const
+    for (const code of cases) {
+      for (let end = 1; end <= code.length; end++) {
+        const prefix = code.slice(0, end)
+        for (const language of ['typescript', 'markup'] as const) {
+          const html = Prism.highlight(prefix, Prism.languages[language], language)
+          const before = document.createElement('pre')
+          before.innerHTML = html
+          const after = document.createElement('pre')
+          for (const line of splitHighlightedLines(html)) {
+            const span = document.createElement('span')
+            span.innerHTML = line
+            after.appendChild(span)
+          }
+          expect(after.textContent).toBe(before.textContent)
+          expect(styledCharacters(after)).toEqual(styledCharacters(before))
+          expect(after.querySelector('script, img')).toBeNull()
+        }
+      }
+    }
+  })
+
+  it('retains unchanged token nodes and updates earlier lines when grammar changes', async () => {
+    ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
+    const container = document.createElement('pre')
+    const root = createRoot(container)
+    const render = (code: string) =>
+      act(async () => {
+        root.render(
+          <HighlightedLines
+            html={Prism.highlight(code, Prism.languages.typescript, 'typescript')}
+          />
+        )
+      })
+    try {
+      await render('const first = 1\nconst second = 2')
+      const firstToken = container.querySelector('.token.keyword')
+      await render('const first = 1\nconst second = 234\nconst third = 3')
+      expect(container.querySelector('.token.keyword')).toBe(firstToken)
+      await render('/* first line\nsecond line')
+      await render('/* first line\nsecond line */\nconst end = 3')
+      expect(container.querySelector('.token.comment')?.textContent).toBe('/* first line')
+      expect(container.textContent).toBe('/* first line\nsecond line */\nconst end = 3')
+    } finally {
+      await act(async () => root.unmount())
+    }
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines.tsx

@@ -0,0 +1,47 @@
+import { memo } from 'react'
+
+/**
+ * Partition Prism output after highlighting the entire fence, so multiline
+ * grammar state stays intact. Reopen crossing spans on each line: React can
+ * retain unchanged line HTML instead of replacing every token in the fence.
+ * This accepts only Prism's generated spans, never raw code or arbitrary HTML.
+ */
+export function splitHighlightedLines(html: string): string[] {
+  const lines: string[] = []
+  const openSpans: string[] = []
+  let line = ''
+  let offset = 0
+  for (const match of html.matchAll(/<span\b[^>]*>|<\/span>|\n/g)) {
+    line += html.slice(offset, match.index)
+    const token = match[0]
+    offset = match.index + token.length
+    if (token === '\n') {
+      lines.push(`${line}${'</span>'.repeat(openSpans.length)}\n`)
+      line = openSpans.join('')
+    } else {
+      if (token === '</span>') openSpans.pop()
+      else openSpans.push(token)
+      line += token
+    }
+  }
+  lines.push(line + html.slice(offset))
+  return lines
+}
+
+interface HighlightedLineProps {
+  html: string
+}
+
+const HighlightedLine = memo(function HighlightedLine({ html }: HighlightedLineProps) {
+  return <span dangerouslySetInnerHTML={{ __html: html }} />
+})
+
+interface HighlightedLinesProps {
+  html: string
+}
+
+export function HighlightedLines({ html }: HighlightedLinesProps) {
+  return splitHighlightedLines(html).map((line, index) => (
+    <HighlightedLine key={index} html={line} />
+  ))
+}

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text.test.ts

@@ -0,0 +1,71 @@
+import remarkGfm from 'remark-gfm'
+import remarkParse from 'remark-parse'
+import { unified } from 'unified'
+import { describe, expect, it } from 'vitest'
+import { remarkPlainText } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text'
+
+const reference = unified().use(remarkParse).use(remarkGfm)
+const optimized = unified().use(remarkParse).use(remarkGfm).use(remarkPlainText)
+
+describe('plain-paragraph parsing', () => {
+  it('matches remark positions, whitespace, and Unicode text', () => {
+    for (const text of [
+      'Hello world! ',
+      'Hello  ',
+      'A café, 日本語, العربية, 😀 and nonbreaking\u00a0spaces.\u00a0',
+      'A sentence with (parentheses), quotes, 1.23 and 20% off.',
+      '12)word',
+      'Plain\u000btext\u000ctext',
+    ] as const) {
+      for (let end = 0; end <= text.length; end++) {
+        const prefix = text.slice(0, end)
+        expect(optimized.parse(prefix), JSON.stringify(prefix)).toEqual(reference.parse(prefix))
+      }
+    }
+  })
+
+  it('preserves Markdown interpretation when a plain prefix gains syntax', () => {
+    const cases = [
+      'Text **bold**, _italic_, ~~deleted~~ and `code`.',
+      'Text [link](https://example.com), ![image](files/chart.png).',
+      'Text www.example.com, https://example.com and user@example.com.',
+      'Text WWW.EXAMPLE.COM and foo@bar.test.',
+      'Text &amp; and &#x41; and <b>raw</b>.',
+      'Text \\*escaped\\* and \\\nline break.',
+      'Text  \nline break.\r\n\r\nNext paragraph.',
+      'Text\u0000replacement\tand tab',
+      '1. Ordered list\n2. Next item',
+      '1) Ordered list',
+      '   Indented text\n\n    code',
+      '# Heading\n\n> Quote\n\n---\n\n- [x] task',
+      'Heading\n=======\n\n| Table |\n| --- |\n| Cell |',
+      '[link][ref]\n\n[ref]: https://example.com',
+      '```typescript\nconst text = `multiline\nvalue`\n```',
+      '\ufeffText',
+    ] as const
+    for (const text of cases) {
+      for (let end = 0; end <= text.length; end++) {
+        const prefix = text.slice(0, end)
+        expect(optimized.parse(prefix), JSON.stringify(prefix)).toEqual(reference.parse(prefix))
+      }
+    }
+  })
+
+  it('matches the reference on growing paragraphs and every ASCII character', () => {
+    const cases = [
+      'A plain paragraph describing the workflow. '.repeat(50),
+      '## Heading\n\nA paragraph.\n\n- A list item\n\n| Column |\n| --- |\n| Cell |\n\n'.repeat(30),
+      `\`\`\`typescript\n${'const value = items.map((item) => item.value)\n'.repeat(40)}\`\`\``,
+      'A citation. <source>{"url":"https://example.com"}</source>\n\n'.repeat(30),
+    ] as const
+    for (const text of cases) {
+      for (let end = 7; end <= text.length; end += 7) {
+        expect(optimized.parse(text.slice(0, end))).toEqual(reference.parse(text.slice(0, end)))
+      }
+    }
+    for (let char = 0; char < 128; char++) {
+      const text = `Before ${String.fromCharCode(char)} after `
+      expect(optimized.parse(text), JSON.stringify(text)).toEqual(reference.parse(text))
+    }
+  })
+})

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text.ts

@@ -0,0 +1,39 @@
+import type { Root } from 'mdast'
+import type { Plugin } from 'unified'
+
+/**
+ * A conservative shortcut for a single plain paragraph. Everything that could
+ * introduce Markdown, GFM autolinks, escapes, entities, or another block still
+ * goes through remark-parse. This changes only tokenization; the usual rehype,
+ * sanitization, components, and word-fade pipeline receive the same tree.
+ */
+export const remarkPlainText: Plugin<[], Root> = function remarkPlainText() {
+  const parse = this.parser
+  if (!parse) throw new Error('remarkPlainText requires remark-parse')
+  this.parser = (document, file) => {
+    if (
+      !/^[A-Za-z0-9]/.test(document) ||
+      /[\n\r\t\0\\`*_{}[\]<>~&#@:|]/.test(document) ||
+      /^\d+[.)](?: |$)/.test(document) ||
+      /www\./i.test(document)
+    ) {
+      return parse(document, file)
+    }
+    const value = document.replace(/ +$/, '')
+    const position = (end: number) => ({
+      start: { line: 1, column: 1, offset: 0 },
+      end: { line: 1, column: end + 1, offset: end },
+    })
+    return {
+      type: 'root',
+      children: [
+        {
+          type: 'paragraph',
+          children: [{ type: 'text', value, position: position(value.length) }],
+          position: position(document.length),
+        },
+      ],
+      position: position(document.length),
+    }
+  }
+}