Skip to content

Commit 33e83ba

Browse files
committed
improvement(chat): reduce streaming render work and completion delay
1 parent b4b2625 commit 33e83ba

7 files changed

Lines changed: 374 additions & 10 deletions

File tree

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import {
1111
useState,
1212
} from 'react'
1313
import type { Nodes } from 'hast'
14-
import { defaultRehypePlugins, Streamdown } from 'streamdown'
14+
import { defaultRehypePlugins, defaultRemarkPlugins, Streamdown } from 'streamdown'
1515
import 'streamdown/styles.css'
1616
// prismjs core must load before its language components — they register on the
1717
// global `Prism` it installs (on `window`/`global`); fixes SSR + client order.
@@ -35,6 +35,8 @@ import {
3535
isInlineFileReference,
3636
} from '@/lib/mothership/chat/inline-image-reference'
3737
import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
38+
import { HighlightedLines } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines'
39+
import { remarkPlainText } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text'
3840
import {
3941
SourceChip,
4042
sourceLabel,
@@ -66,6 +68,8 @@ const LANG_ALIASES: Record<string, string> = {
6668
py: 'python',
6769
}
6870

71+
const MARKDOWN_REMARK_PLUGINS = [...Object.values(defaultRemarkPlugins), remarkPlainText]
72+
6973
const PROSE_CLASSES = cn(
7074
'prose prose-base dark:prose-invert max-w-none',
7175
'font-[family-name:var(--font-inter)] antialiased break-words tracking-[0]',
@@ -308,10 +312,9 @@ const MARKDOWN_COMPONENTS = {
308312
/>
309313
</div>
310314
<div className='code-editor-theme bg-[var(--surface-5)] dark:bg-[var(--code-bg)]'>
311-
<pre
312-
className='m-0 overflow-x-auto whitespace-pre p-4 font-mono text-[var(--text-primary)] text-small leading-[21px]'
313-
dangerouslySetInnerHTML={{ __html: html }}
314-
/>
315+
<pre className='m-0 overflow-x-auto whitespace-pre p-4 font-mono text-[var(--text-primary)] text-small leading-[21px]'>
316+
<HighlightedLines html={html} />
317+
</pre>
315318
</div>
316319
</div>
317320
)
@@ -704,6 +707,7 @@ function ChatContentInner({
704707
isAnimating={streamingTree}
705708
components={MARKDOWN_COMPONENTS}
706709
rehypePlugins={imageRehypePlugins}
710+
remarkPlugins={MARKDOWN_REMARK_PLUGINS}
707711
>
708712
{group.markdown}
709713
</Streamdown>
Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
/** @vitest-environment jsdom */
2+
import { act } from 'react'
3+
import Prism from 'prismjs'
4+
import { createRoot } from 'react-dom/client'
5+
import 'prismjs/components/prism-typescript'
6+
import { describe, expect, it } from 'vitest'
7+
import {
8+
HighlightedLines,
9+
splitHighlightedLines,
10+
} from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/highlighted-lines'
11+
12+
/** Compare per-character token ancestry, ignoring unstyled line wrappers. */
13+
function styledCharacters(element: Element): Array<[string, string]> {
14+
const result: Array<[string, string]> = []
15+
function visit(node: Node, classes: string[]) {
16+
if (node.nodeType === Node.TEXT_NODE) {
17+
for (const char of node.textContent ?? '') {
18+
if (char !== '\n') result.push([char, classes.join('/')])
19+
}
20+
} else if (node instanceof Element) {
21+
const nextClasses = node.className ? [...classes, node.className] : classes
22+
for (const child of node.childNodes) visit(child, nextClasses)
23+
}
24+
}
25+
visit(element, [])
26+
return result
27+
}
28+
29+
describe('highlighted code lines', () => {
30+
it('preserves full-fence Prism text and token styles across multiline constructs', () => {
31+
const cases = [
32+
`const value = \`first\nsecond \${1 + 2}\nthird\`\n`,
33+
'/* a comment\nwith <script>alert(1)</script>\n*/\nconst x = "safe"',
34+
'<div title="a &amp; b">\n <!-- multi\n line -->\n</div>',
35+
'const x = "<img src=x onerror=alert(1)>"\n\nconst y = 1',
36+
]
37+
for (const code of cases) {
38+
for (let end = 1; end <= code.length; end++) {
39+
const prefix = code.slice(0, end)
40+
for (const language of ['typescript', 'markup']) {
41+
const html = Prism.highlight(prefix, Prism.languages[language], language)
42+
const before = document.createElement('pre')
43+
before.innerHTML = html
44+
const after = document.createElement('pre')
45+
for (const line of splitHighlightedLines(html)) {
46+
const span = document.createElement('span')
47+
span.innerHTML = line
48+
after.appendChild(span)
49+
}
50+
expect(after.textContent).toBe(before.textContent)
51+
expect(styledCharacters(after)).toEqual(styledCharacters(before))
52+
expect(after.querySelector('script, img')).toBeNull()
53+
}
54+
}
55+
}
56+
})
57+
58+
it('retains unchanged token nodes and updates earlier lines when grammar changes', async () => {
59+
;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
60+
const container = document.createElement('pre')
61+
const root = createRoot(container)
62+
const render = (code: string) =>
63+
act(async () => {
64+
root.render(
65+
<HighlightedLines
66+
html={Prism.highlight(code, Prism.languages.typescript, 'typescript')}
67+
/>
68+
)
69+
})
70+
try {
71+
await render('const first = 1\nconst second = 2')
72+
const firstToken = container.querySelector('.token.keyword')
73+
await render('const first = 1\nconst second = 234\nconst third = 3')
74+
expect(container.querySelector('.token.keyword')).toBe(firstToken)
75+
await render('/* first line\nsecond line')
76+
await render('/* first line\nsecond line */\nconst end = 3')
77+
expect(container.querySelector('.token.comment')?.textContent).toBe('/* first line')
78+
expect(container.textContent).toBe('/* first line\nsecond line */\nconst end = 3')
79+
} finally {
80+
await act(async () => root.unmount())
81+
}
82+
})
83+
})
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import { memo } from 'react'
2+
3+
/**
4+
* Partition Prism output after highlighting the entire fence, so multiline
5+
* grammar state stays intact. Reopen crossing spans on each line: React can
6+
* retain unchanged line HTML instead of replacing every token in the fence.
7+
* This accepts only Prism's generated spans, never raw code or arbitrary HTML.
8+
*/
9+
export function splitHighlightedLines(html: string): string[] {
10+
const lines: string[] = []
11+
const openSpans: string[] = []
12+
let line = ''
13+
let offset = 0
14+
for (const match of html.matchAll(/<span\b[^>]*>|<\/span>|\n/g)) {
15+
line += html.slice(offset, match.index)
16+
const token = match[0]
17+
offset = match.index + token.length
18+
if (token === '\n') {
19+
lines.push(`${line}${'</span>'.repeat(openSpans.length)}\n`)
20+
line = openSpans.join('')
21+
} else {
22+
if (token === '</span>') openSpans.pop()
23+
else openSpans.push(token)
24+
line += token
25+
}
26+
}
27+
lines.push(line + html.slice(offset))
28+
return lines
29+
}
30+
31+
interface HighlightedLineProps {
32+
html: string
33+
}
34+
35+
const HighlightedLine = memo(function HighlightedLine({ html }: HighlightedLineProps) {
36+
return <span dangerouslySetInnerHTML={{ __html: html }} />
37+
})
38+
39+
interface HighlightedLinesProps {
40+
html: string
41+
}
42+
43+
export function HighlightedLines({ html }: HighlightedLinesProps) {
44+
return splitHighlightedLines(html).map((line, index) => (
45+
<HighlightedLine key={index} html={line} />
46+
))
47+
}
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import remarkGfm from 'remark-gfm'
2+
import remarkParse from 'remark-parse'
3+
import { unified } from 'unified'
4+
import { describe, expect, it } from 'vitest'
5+
import { remarkPlainText } from '@/app/workspace/[workspaceId]/home/components/message-content/components/chat-content/remark-plain-text'
6+
7+
const reference = unified().use(remarkParse).use(remarkGfm)
8+
const optimized = unified().use(remarkParse).use(remarkGfm).use(remarkPlainText)
9+
10+
describe('plain-paragraph parsing', () => {
11+
it('matches remark positions, whitespace, and Unicode text', () => {
12+
for (const text of [
13+
'Hello world! ',
14+
'Hello ',
15+
'A café, 日本語, العربية, 😀 and nonbreaking\u00a0spaces.\u00a0',
16+
'A sentence with (parentheses), quotes, 1.23 and 20% off.',
17+
'12)word',
18+
'Plain\u000btext\u000ctext',
19+
]) {
20+
for (let end = 0; end <= text.length; end++) {
21+
const prefix = text.slice(0, end)
22+
expect(optimized.parse(prefix), JSON.stringify(prefix)).toEqual(reference.parse(prefix))
23+
}
24+
}
25+
})
26+
27+
it('preserves Markdown interpretation when a plain prefix gains syntax', () => {
28+
const cases = [
29+
'Text **bold**, _italic_, ~~deleted~~ and `code`.',
30+
'Text [link](https://example.com), ![image](files/chart.png).',
31+
'Text www.example.com, https://example.com and user@example.com.',
32+
'Text WWW.EXAMPLE.COM and foo@bar.test.',
33+
'Text &amp; and &#x41; and <b>raw</b>.',
34+
'Text \\*escaped\\* and \\\nline break.',
35+
'Text \nline break.\r\n\r\nNext paragraph.',
36+
'Text\u0000replacement\tand tab',
37+
'1. Ordered list\n2. Next item',
38+
'1) Ordered list',
39+
' Indented text\n\n code',
40+
'# Heading\n\n> Quote\n\n---\n\n- [x] task',
41+
'Heading\n=======\n\n| Table |\n| --- |\n| Cell |',
42+
'[link][ref]\n\n[ref]: https://example.com',
43+
'```typescript\nconst text = `multiline\nvalue`\n```',
44+
'\ufeffText',
45+
]
46+
for (const text of cases) {
47+
for (let end = 0; end <= text.length; end++) {
48+
const prefix = text.slice(0, end)
49+
expect(optimized.parse(prefix), JSON.stringify(prefix)).toEqual(reference.parse(prefix))
50+
}
51+
}
52+
})
53+
54+
it('matches the reference on growing paragraphs and every ASCII character', () => {
55+
const cases = [
56+
'A plain paragraph describing the workflow. '.repeat(50),
57+
'## Heading\n\nA paragraph.\n\n- A list item\n\n| Column |\n| --- |\n| Cell |\n\n'.repeat(30),
58+
`\`\`\`typescript\n${'const value = items.map((item) => item.value)\n'.repeat(40)}\`\`\``,
59+
'A citation. <source>{"url":"https://example.com"}</source>\n\n'.repeat(30),
60+
]
61+
for (const text of cases) {
62+
for (let end = 7; end <= text.length; end += 7) {
63+
expect(optimized.parse(text.slice(0, end))).toEqual(reference.parse(text.slice(0, end)))
64+
}
65+
}
66+
for (let char = 0; char < 128; char++) {
67+
const text = `Before ${String.fromCharCode(char)} after `
68+
expect(optimized.parse(text), JSON.stringify(text)).toEqual(reference.parse(text))
69+
}
70+
})
71+
})
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import type { Root } from 'mdast'
2+
import type { Plugin } from 'unified'
3+
4+
/**
5+
* A conservative shortcut for a single plain paragraph. Everything that could
6+
* introduce Markdown, GFM autolinks, escapes, entities, or another block still
7+
* goes through remark-parse. This changes only tokenization; the usual rehype,
8+
* sanitization, components, and word-fade pipeline receive the same tree.
9+
*/
10+
export const remarkPlainText: Plugin<[], Root> = function remarkPlainText() {
11+
const parse = this.parser
12+
if (!parse) throw new Error('remarkPlainText requires remark-parse')
13+
this.parser = (document, file) => {
14+
if (
15+
!/^[A-Za-z0-9]/.test(document) ||
16+
/[\n\r\t\0\\`*_{}[\]<>~&#@:|]/.test(document) ||
17+
/^\d+[.)](?: |$)/.test(document) ||
18+
/www\./i.test(document)
19+
) {
20+
return parse(document, file)
21+
}
22+
const value = document.replace(/ +$/, '')
23+
const position = (end: number) => ({
24+
start: { line: 1, column: 1, offset: 0 },
25+
end: { line: 1, column: end + 1, offset: end },
26+
})
27+
return {
28+
type: 'root',
29+
children: [
30+
{
31+
type: 'paragraph',
32+
children: [{ type: 'text', value, position: position(value.length) }],
33+
position: position(document.length),
34+
},
35+
],
36+
position: position(document.length),
37+
}
38+
}
39+
}

0 commit comments

Comments
 (0)