perf(desktop): render streaming replies incrementally

CHANGELOG.md

@@ -35,6 +35,7 @@ All notable changes to DreamCoder will be documented in this file.
 
 - **Sidebar re-renders** (#157, #158): the sidebar no longer re-renders on streamed text updates. Running indicators still follow session and tab status. Thanks @tedy69.
 - **Re-renders while a reply runs** (#156, #165): the session page no longer re-renders every second during a reply. Only the elapsed time updates, and it now counts from when the server starts the turn. Thanks @princecherry9905.
+- **Smoother streaming replies** (#154): a reply that is still streaming is no longer re-parsed and redrawn in full on every update. Finished paragraphs, lists, tables and code blocks are drawn once and stay in place; only the part still being written is redrawn. Text selected in the finished part stays selected while the reply continues. Code in a streaming reply uses the lighter highlighter and gets full highlighting when the reply ends, so code blocks no longer flicker. Following the bottom of the reply now scrolls at most once per frame.
 - New sessions start on the active provider's main model. Before, activating a provider saved its main model in a hidden setting, and that saved value kept overriding the provider after its main model was edited.
 - The selected model card in the dark theme had light pink colors with text that was hard to read.
 

desktop/src/components/chat/AssistantMessage.test.tsx

@@ -0,0 +1,27 @@
+import { describe, expect, it, vi } from 'vitest'
+import { render, screen } from '@testing-library/react'
+import '@testing-library/jest-dom'
+
+vi.mock('./CodeViewer', () => ({
+  CodeViewer: ({ code, highlighter }: { code: string; highlighter?: string }) => (
+    <div data-testid="code-viewer" data-highlighter={highlighter ?? 'shiki'}>{code}</div>
+  ),
+}))
+
+import { AssistantMessage } from './AssistantMessage'
+
+const REPLY = 'Here is the code:\n\n```ts\nconst a = 1\n```\n\nDone.'
+
+describe('AssistantMessage', () => {
+  it('renders a streaming reply incrementally with the lightweight highlighter', () => {
+    const { container } = render(<AssistantMessage content={REPLY} isStreaming incremental />)
+    expect(container.querySelector('[data-streaming-block]')).not.toBeNull()
+    expect(screen.getByTestId('code-viewer')).toHaveAttribute('data-highlighter', 'prism')
+  })
+
+  it('renders a committed reply in full with Shiki', () => {
+    const { container } = render(<AssistantMessage content={REPLY} />)
+    expect(container.querySelector('[data-streaming-block]')).toBeNull()
+    expect(screen.getByTestId('code-viewer')).toHaveAttribute('data-highlighter', 'shiki')
+  })
+})

desktop/src/components/chat/AssistantMessage.tsx

@@ -1,17 +1,18 @@
 import { memo } from 'react'
-import { MarkdownRenderer } from '../markdown/MarkdownRenderer'
+import { MarkdownRenderer, StreamingMarkdownRenderer } from '../markdown/MarkdownRenderer'
 import { MessageActionBar, type MessageBranchAction } from './MessageActionBar'
 import { InlineImageGallery } from './InlineImageGallery'
 
 type Props = {
   content: string
   isStreaming?: boolean
+  incremental?: boolean
   branchAction?: MessageBranchAction
   /** The reply's actions, shown under its last text: copy takes the whole reply. Omitted on earlier texts. */
   replyCopyText?: string
 }
 
-export const AssistantMessage = memo(function AssistantMessage({ content, isStreaming, branchAction, replyCopyText }: Props) {
+export const AssistantMessage = memo(function AssistantMessage({ content, isStreaming, incremental, branchAction, replyCopyText }: Props) {
   if (!content.trim()) return null
 
   const documentLayout = shouldUseDocumentLayout(content)
@@ -24,12 +25,19 @@ export const AssistantMessage = memo(function AssistantMessage({ content, isStre
         className="flex w-full min-w-0 max-w-full flex-col items-start gap-1.5"
       >
         <div className="w-full min-w-0 text-ui-message leading-[1.75] text-[var(--color-text-primary)]">
-          <MarkdownRenderer
-            content={content}
-            variant={documentLayout ? 'document' : 'default'}
-            streaming={isStreaming}
-            className={MESSAGE_PROSE_CLASSES}
-          />
+          {incremental ? (
+            <StreamingMarkdownRenderer
+              content={content}
+              variant={documentLayout ? 'document' : 'default'}
+              className={MESSAGE_PROSE_CLASSES}
+            />
+          ) : (
+            <MarkdownRenderer
+              content={content}
+              variant={documentLayout ? 'document' : 'default'}
+              className={MESSAGE_PROSE_CLASSES}
+            />
+          )}
           {!isStreaming && <InlineImageGallery text={content} />}
           {isStreaming && (
             <span className="ml-0.5 inline-block h-4 w-0.5 animate-shimmer bg-[var(--color-brand)] align-text-bottom" />

desktop/src/components/chat/CodeViewer.test.tsx

@@ -1,7 +1,19 @@
-import { render, screen } from '@testing-library/react'
-import { describe, expect, it } from 'vitest'
+import { render, screen, waitFor } from '@testing-library/react'
+import { afterEach, describe, expect, it, vi } from 'vitest'
 import { CodeViewer } from './CodeViewer'
 
+const shiki = vi.hoisted(() => ({ loads: 0 }))
+
+vi.mock('react-shiki', () => {
+  shiki.loads += 1
+  return {
+    ShikiHighlighter: ({ children }: { children: string }) => (
+      <div data-testid="shiki-container"><code>{children}</code></div>
+    ),
+    createJavaScriptRegexEngine: () => ({}),
+  }
+})
+
 describe('CodeViewer', () => {
   it('keeps the same inner padding for highlighted code content', () => {
     const { container } = render(
@@ -22,4 +34,30 @@ describe('CodeViewer', () => {
     expect(container.querySelector('[data-line-number="1"]')).toBeTruthy()
     expect(container.querySelector('[data-line-number="2"]')).toBeTruthy()
   })
+
+  describe('highlighter', () => {
+    afterEach(() => {
+      vi.unstubAllEnvs()
+    })
+
+    it('never loads Shiki for the prism highlighter', async () => {
+      vi.stubEnv('MODE', 'production')
+      const { container } = render(<CodeViewer code="const a = 1" language="ts" highlighter="prism" />)
+
+      await new Promise((resolve) => setTimeout(resolve, 20))
+      expect(shiki.loads).toBe(0)
+      expect(container.querySelector('[data-highlight-engine="shiki"]')).toBeNull()
+      expect(container.querySelector('[data-highlight-engine="prism"]')).not.toBeNull()
+    })
+
+    it('switches to Shiki for the default highlighter', async () => {
+      vi.stubEnv('MODE', 'production')
+      const { container } = render(<CodeViewer code="const b = 2" language="ts" />)
+
+      await waitFor(() => {
+        expect(container.querySelector('[data-highlight-engine="shiki"] [data-testid="shiki-container"]')).not.toBeNull()
+      })
+      expect(shiki.loads).toBe(1)
+    })
+  })
 })

desktop/src/components/chat/CodeViewer.tsx

@@ -8,6 +8,7 @@ type Props = {
   language?: string
   maxLines?: number
   showLineNumbers?: boolean
+  highlighter?: 'shiki' | 'prism'
 }
 
 const warmPrismTheme: PrismTheme = {
@@ -169,12 +170,13 @@ function PrismCodeContent({ code, language, showLineNumbers }: { code: string; l
   )
 }
 
-function CodeArea({ code, language, showLineNumbers }: { code: string; language?: string; showLineNumbers: boolean }) {
+function CodeArea({ code, language, showLineNumbers, highlighter }: { code: string; language?: string; showLineNumbers: boolean; highlighter: 'shiki' | 'prism' }) {
   const containerRef = useRef<HTMLDivElement>(null)
   const [runtime, setRuntime] = useState<ShikiRuntime | null>(null)
   const [loaded, setLoaded] = useState(false)
 
   useEffect(() => {
+    if (highlighter === 'prism') return
     let cancelled = false
     setLoaded(false)
     loadShikiRuntime().then((nextRuntime) => {
@@ -183,7 +185,7 @@ function CodeArea({ code, language, showLineNumbers }: { code: string; language?
     return () => {
       cancelled = true
     }
-  }, [])
+  }, [highlighter])
 
   useEffect(() => {
     setLoaded(false)
@@ -205,7 +207,8 @@ function CodeArea({ code, language, showLineNumbers }: { code: string; language?
     return () => observer.disconnect()
   }, [runtime, code, language])
 
-  const ShikiHighlighter = runtime?.Highlighter
+  const shikiRuntime = highlighter === 'shiki' ? runtime : null
+  const ShikiHighlighter = shikiRuntime?.Highlighter
 
   return (
     <div
@@ -239,7 +242,7 @@ function CodeArea({ code, language, showLineNumbers }: { code: string; language?
           <ShikiHighlighter
             language={language || 'text'}
             theme={warmShikiTheme}
-            engine={runtime.engine}
+            engine={shikiRuntime.engine}
             showLineNumbers={showLineNumbers}
             showLanguage={false}
             addDefaultStyles={false}
@@ -259,7 +262,7 @@ function CodeArea({ code, language, showLineNumbers }: { code: string; language?
   )
 }
 
-export function CodeViewer({ code, language, maxLines = 20, showLineNumbers = false }: Props) {
+export function CodeViewer({ code, language, maxLines = 20, showLineNumbers = false, highlighter = 'shiki' }: Props) {
   const [expanded, setExpanded] = useState(false)
   const t = useTranslation()
 
@@ -289,6 +292,7 @@ export function CodeViewer({ code, language, maxLines = 20, showLineNumbers = fa
         code={visibleCode}
         language={language}
         showLineNumbers={effectiveShowLineNumbers}
+        highlighter={highlighter}
       />
 
       {showExpandToggle && (

desktop/src/components/chat/MessageList.test.tsx

@@ -1847,6 +1847,7 @@ describe('MessageList nested tool calls', () => {
     await waitFor(() => {
       expect(screen.getByText('streaming next token')).toBeTruthy()
     })
+    await waitForProgrammaticScrollReset()
     expect(scrollIntoView).not.toHaveBeenCalled()
     expect(scrollTop).toBe(600)
   })
@@ -1985,6 +1986,7 @@ describe('MessageList nested tool calls', () => {
     await waitFor(() => {
       expect(screen.getByText('streaming next token after height change')).toBeTruthy()
     })
+    await waitForProgrammaticScrollReset()
     expect(scrollIntoView).not.toHaveBeenCalled()
     expect(scrollTop).toBe(1000)
 
@@ -2059,6 +2061,7 @@ describe('MessageList nested tool calls', () => {
     act(() => {
       resizeCallback?.([], {} as ResizeObserver)
     })
+    await waitForProgrammaticScrollReset()
 
     expect(scrollIntoView).not.toHaveBeenCalled()
     expect(scrollTop).toBe(1200)
@@ -2132,12 +2135,14 @@ describe('MessageList nested tool calls', () => {
     act(() => {
       resizeCallback?.(makeResizeEntry(400), {} as ResizeObserver)
     })
+    await waitForProgrammaticScrollReset()
     expect(scrollTop).toBe(600)
 
     scrollTopWriteCount = 0
     act(() => {
       resizeCallback?.(makeResizeEntry(401), {} as ResizeObserver)
     })
+    await waitForProgrammaticScrollReset()
     act(() => {
       resizeCallback?.(makeResizeEntry(400), {} as ResizeObserver)
     })
@@ -2149,6 +2154,7 @@ describe('MessageList nested tool calls', () => {
     act(() => {
       resizeCallback?.(makeResizeEntry(420), {} as ResizeObserver)
     })
+    await waitForProgrammaticScrollReset()
 
     expect(scrollTop).toBe(640)
   })