fix(desktop): reduce long-history streaming work

apps/desktop/scripts/perf/README.md

@@ -19,6 +19,14 @@ npm run perf                  # attaches, runs the CI suite, gates on baseline
 # One scenario, with a CPU profile:
 npm run perf -- stream --cpuprofile --tokens 800
 
+# Report-only stream over a preloaded 200-turn transcript. The harness verifies
+# the source count, waits two paints plus a configurable settle delay, then starts
+# recorders; slower machines may still include residual mount/highlight work:
+npm run perf -- stream-history --spawn --prod
+
+# Stronger local stress point (1000 turns / 2000 messages):
+npm run perf -- stream-history --spawn --prod --historyTurns 1000
+
 # Representative PRODUCTION numbers (minified React, not the ~3x-slower dev build):
 npm run perf -- cold-start stream keystroke transcript --spawn --prod
 

apps/desktop/scripts/perf/scenarios/index.mjs

@@ -7,11 +7,13 @@ import keystroke from './keystroke.mjs'
 import profileSwitch from './profile-switch.mjs'
 import sessionSwitch from './session-switch.mjs'
 import stream from './stream.mjs'
+import streamHistory from './stream-history.mjs'
 import submit from './submit.mjs'
 import transcript from './transcript.mjs'
 
 export const SCENARIOS = {
   [stream.name]: stream,
+  [streamHistory.name]: streamHistory,
   [keystroke.name]: keystroke,
   [transcript.name]: transcript,
   [coldStart.name]: coldStart,

apps/desktop/scripts/perf/scenarios/stream-history.mjs

@@ -0,0 +1,13 @@
+import stream from './stream.mjs'
+
+export default {
+  name: 'stream-history',
+  tier: 'manual',
+  description: 'Synthetic token stream over a preloaded long transcript.',
+  run(cdp, opts = {}) {
+    return stream.run(cdp, {
+      ...opts,
+      historyTurns: opts.historyTurns ?? 200
+    })
+  }
+}

apps/desktop/scripts/perf/scenarios/stream.mjs

@@ -119,6 +119,8 @@ export default {
     const tokens = Number(opts.tokens ?? 400)
     const intervalMs = Number(opts.intervalMs ?? 16)
     const flushMinMs = Number(opts.flushMinMs ?? 33)
+    const historyTurns = Number(opts.historyTurns ?? 0)
+    const historySettleMs = Number(opts.historySettleMs ?? 1500)
     // Realistic default: a short markdown paragraph ending in a blank line, so
     // blocks SETTLE as they stream — exactly how real LLM output behaves, and
     // what block-memoization is designed for (only the growing tail re-renders).
@@ -130,6 +132,25 @@ export default {
     const real = Boolean(opts.real)
 
     await cdp.send('Runtime.enable')
+
+    if (historyTurns > 0) {
+      if (real) {
+        throw new Error('--historyTurns is only supported by the synthetic stream scenario')
+      }
+
+      // Preload history before recorders start, then allow a configurable
+      // delay. Source count is verified below; slower renderers may still have
+      // residual mount/highlight work when recording begins.
+      await cdp.eval(`window.__PERF_DRIVE__.loadTranscript(${historyTurns})`)
+      await sleep(historySettleMs)
+      const mountedMessages = Number(await cdp.eval('window.__PERF_DRIVE__.snapshotMsgs()'))
+      const expectedMessages = historyTurns * 2
+
+      if (mountedMessages !== expectedMessages) {
+        throw new Error(`expected ${expectedMessages} preloaded history messages, got ${mountedMessages}`)
+      }
+    }
+
     await cdp.eval(RECORDERS)
 
     if (real) {
@@ -177,7 +198,19 @@ export default {
       )
       await sleep(200)
       await cdp.eval('window.__MO__.arm()')
-      await sleep(tokens * intervalMs + 1500)
+      const syntheticTimeoutMs = Number(opts.timeoutMs ?? Math.max(tokens * intervalMs + 10000, 120000))
+      const deadline = Date.now() + syntheticTimeoutMs
+
+      while (Date.now() < deadline && (await cdp.eval('window.__PERF_DRIVE__.streaming()'))) {
+        await sleep(50)
+      }
+
+      if (await cdp.eval('window.__PERF_DRIVE__.streaming()')) {
+        throw new Error('synthetic stream did not finish before the timeout')
+      }
+
+      // Include the final commit and paint after the driver reports idle.
+      await sleep(250)
     }
 
     const data = JSON.parse(await cdp.eval(COLLECT))
@@ -186,6 +219,9 @@ export default {
       await cdp.eval('window.__PERF_DRIVE__.reset()')
     }
 
-    return analyze(data, real ? 0 : 500)
+    const result = analyze(data, real ? 0 : 500)
+    result.detail.historyTurns = historyTurns
+
+    return result
   }
 }

apps/desktop/src/app/chat/perf-probe.tsx

@@ -1,5 +1,6 @@
 import { Profiler, type ProfilerOnRenderCallback, type ReactNode } from 'react'
 
+import { updateChatMessageAt } from '@/lib/chat-messages'
 import { $gateway } from '@/store/gateway'
 import { $messages, setBusy, setMessages } from '@/store/session'
 
@@ -24,7 +25,13 @@ declare global {
     }
     __PERF_DRIVE__?: {
       /** Inject an assistant message and grow it by `chunk` every `intervalMs`. Returns a stop handle. */
-      stream: (opts?: { chunk?: string; intervalMs?: number; totalTokens?: number }) => SyntheticDriverHandle
+      stream: (opts?: {
+        chunk?: string
+        intervalMs?: number
+        totalTokens?: number
+        flushMinMs?: number
+      }) => SyntheticDriverHandle
+      streaming: () => boolean
       /**
        * Replace the transcript with `turns` synthetic user/assistant pairs of
        * realistic mixed markdown, then resolve with the ms elapsed from the
@@ -159,6 +166,7 @@ if (typeof window !== 'undefined' && !window.__PERF_DRIVE__) {
 
   window.__PERF_DRIVE__ = {
     snapshotMsgs: () => $messages.get().length,
+    streaming: () => activeHandle !== null,
     connected: () => {
       try {
         return $gateway.get()?.connectionState === 'open'
@@ -242,22 +250,20 @@ if (typeof window !== 'undefined' && !window.__PERF_DRIVE__) {
           return
         }
 
-        setMessages(prev =>
-          prev.map(m => {
-            if (m.id !== msgId) {
-              return m
-            }
+        setMessages(prev => {
+          const index = prev.at(-1)?.id === msgId ? prev.length - 1 : prev.findIndex(message => message.id === msgId)
 
-            const head = m.parts.slice(0, -1)
-            const last = m.parts.at(-1)
+          return updateChatMessageAt(prev, index, message => {
+            const head = message.parts.slice(0, -1)
+            const last = message.parts.at(-1)
             const lastText = last && last.type === 'text' ? last.text : ''
 
             return {
-              ...m,
+              ...message,
               parts: [...head, { type: 'text', text: lastText + delta }]
             }
           })
-        )
+        })
       }
 
       const flushNow = () => {
@@ -281,10 +287,7 @@ if (typeof window !== 'undefined' && !window.__PERF_DRIVE__) {
 
         const since = performance.now() - lastFlushAt
         const wait = Math.max(0, flushMinMs - since)
-        flushHandle =
-          wait <= 0 && typeof requestAnimationFrame === 'function'
-            ? requestAnimationFrame(flushNow)
-            : (setTimeout(flushNow, wait) as unknown as number)
+        flushHandle = setTimeout(flushNow, wait) as unknown as number
       }
 
       const handle: SyntheticDriverHandle = {
@@ -297,7 +300,6 @@ if (typeof window !== 'undefined' && !window.__PERF_DRIVE__) {
 
           if (flushHandle !== null) {
             clearTimeout(flushHandle)
-            cancelAnimationFrame?.(flushHandle)
           }
 
           flushHandle = null
@@ -309,7 +311,11 @@ if (typeof window !== 'undefined' && !window.__PERF_DRIVE__) {
 
           activeHandle = null
           // Mark message finalized.
-          setMessages(prev => prev.map(m => (m.id === msgId ? { ...m, pending: false } : m)))
+          setMessages(prev => {
+            const index = prev.at(-1)?.id === msgId ? prev.length - 1 : prev.findIndex(message => message.id === msgId)
+
+            return updateChatMessageAt(prev, index, message => ({ ...message, pending: false }))
+          })
           setBusy(false)
         }
       }

apps/desktop/src/app/chat/runtime-repository.test.tsx

@@ -0,0 +1,93 @@
+import { renderHook } from '@testing-library/react'
+import { describe, expect, it } from 'vitest'
+
+import { type ChatMessage, getChatMessageListUpdate, updateChatMessageAt } from '@/lib/chat-messages'
+
+import { runtimeMessageRepositoryDelta, useRuntimeMessageRepository } from './runtime-repository'
+
+function message(id: string, role: ChatMessage['role'], text: string, pending = false): ChatMessage {
+  return {
+    id,
+    role,
+    parts: [{ type: 'text', text }],
+    pending
+  }
+}
+
+describe('useRuntimeMessageRepository', () => {
+  it('requires the exact predecessor when updates are batched', () => {
+    const messages = [message('user-1', 'user', 'Question'), message('assistant-1', 'assistant', 'A', true)]
+    const first = updateChatMessageAt(messages, 1, current => message(current.id, current.role, 'AB', true))
+    const second = updateChatMessageAt(first, 1, current => message(current.id, current.role, 'ABC', true))
+
+    expect(getChatMessageListUpdate(messages, second)).toBeNull()
+    expect(getChatMessageListUpdate(first, second)).toMatchObject({
+      index: 1,
+      previousMessage: first[1],
+      message: second[1]
+    })
+    expect(getChatMessageListUpdate(messages, first)).not.toBeNull()
+  })
+
+  it('preserves settled runtime message identity when only the streaming tail changes', () => {
+    const user = message('user-1', 'user', 'Question')
+    const settled = message('assistant-1', 'assistant', 'Settled answer')
+    const streaming = message('assistant-2', 'assistant', 'A', true)
+
+    const { result, rerender } = renderHook(
+      ({ messages }: { messages: ChatMessage[] }) => useRuntimeMessageRepository(messages),
+      { initialProps: { messages: [user, settled, streaming] } }
+    )
+
+    const initialRepository = result.current
+    const nextStreaming = message('assistant-2', 'assistant', 'AB', true)
+
+    rerender({ messages: [user, settled, nextStreaming] })
+
+    expect(result.current.messages[0]?.message).toBe(initialRepository.messages[0]?.message)
+    expect(result.current.messages[1]?.message).toBe(initialRepository.messages[1]?.message)
+    expect(result.current.messages[2]?.message).not.toBe(initialRepository.messages[2]?.message)
+    expect(result.current.messages.map(item => item.parentId)).toEqual([null, 'user-1', 'assistant-1'])
+    expect(result.current.headId).toBe('assistant-2')
+  })
+
+  it('forwards an annotated pending-tail update without rebuilding settled repository items', () => {
+    const messages = [
+      message('user-1', 'user', 'Question'),
+      message('assistant-1', 'assistant', 'Settled answer'),
+      message('assistant-2', 'assistant', 'A', true)
+    ]
+
+    const { result, rerender } = renderHook(
+      ({ messages: nextMessages }: { messages: ChatMessage[] }) => useRuntimeMessageRepository(nextMessages),
+      { initialProps: { messages } }
+    )
+
+    const initialRepository = result.current
+    const nextMessages = updateChatMessageAt(messages, 2, current => message(current.id, current.role, 'AB', true))
+
+    rerender({ messages: nextMessages })
+
+    expect(result.current.messages).toBe(initialRepository.messages)
+    expect(result.current[runtimeMessageRepositoryDelta]?.message.content).toEqual([{ type: 'text', text: 'AB' }])
+    expect(result.current[runtimeMessageRepositoryDelta]?.parentId).toBe('assistant-1')
+  })
+
+  it('falls back to a complete repository when the tail settles', () => {
+    const messages = [message('user-1', 'user', 'Question'), message('assistant-1', 'assistant', 'A', true)]
+
+    const { result, rerender } = renderHook(
+      ({ messages: nextMessages }: { messages: ChatMessage[] }) => useRuntimeMessageRepository(nextMessages),
+      { initialProps: { messages } }
+    )
+
+    const initialRepository = result.current
+    const nextMessages = updateChatMessageAt(messages, 1, current => message(current.id, current.role, 'Done', false))
+
+    rerender({ messages: nextMessages })
+
+    expect(result.current.messages).not.toBe(initialRepository.messages)
+    expect(result.current[runtimeMessageRepositoryDelta]).toBeUndefined()
+    expect(result.current.messages.at(-1)?.message.status).toEqual({ type: 'complete', reason: 'stop' })
+  })
+})

apps/desktop/src/app/chat/runtime-repository.ts

@@ -1,20 +1,79 @@
-import { ExportedMessageRepository, type ThreadMessage } from '@assistant-ui/react'
+import type { ExportedMessageRepository, ThreadMessage } from '@assistant-ui/react'
 import { useMemo, useRef } from 'react'
 
-import type { ChatMessage } from '@/lib/chat-messages'
+import { type ChatMessage, getChatMessageListUpdate } from '@/lib/chat-messages'
 import { coalesceToolOnlyAssistants, createToolMergeCache, toRuntimeMessage } from '@/lib/chat-runtime'
 
+export const runtimeMessageRepositoryDelta = Symbol('runtimeMessageRepositoryDelta')
+
+export type RuntimeMessageRepository = ExportedMessageRepository & {
+  [runtimeMessageRepositoryDelta]?: {
+    message: ThreadMessage
+    parentId: string | null
+  }
+}
+
+type RepositoryCache = {
+  repository: RuntimeMessageRepository
+  source: ChatMessage[]
+  tailItem: { message: ThreadMessage; parentId: string | null } | null
+  tailSource: ChatMessage | null
+}
+
 /**
  * ChatMessage[] -> assistant-ui message repository, with a WeakMap identity
  * cache so unchanged messages convert once (and a tool-merge cache that folds
  * tool-only assistant turns into their neighbour). Shared by the main chat's
  * runtime boundary and session tiles — one transcript pipeline, N surfaces.
  */
-export function useRuntimeMessageRepository(messages: ChatMessage[]): ExportedMessageRepository {
+export function useRuntimeMessageRepository(messages: ChatMessage[]): RuntimeMessageRepository {
   const cacheRef = useRef(new WeakMap<ChatMessage, ThreadMessage>())
   const toolMergeCacheRef = useRef(createToolMergeCache())
+  const repositoryCacheRef = useRef<RepositoryCache | null>(null)
 
   return useMemo(() => {
+    const previous = repositoryCacheRef.current
+    const update = previous ? getChatMessageListUpdate(previous.source, messages) : null
+    const nextTail = messages.at(-1) ?? null
+
+    // Stream deltas replace only the pending tail. The provenance attached by
+    // updateChatMessageAt lets us forward that one normalized record without
+    // rescanning/coalescing the settled transcript. Completion, branch changes,
+    // hydration, and arbitrary edits fail closed to the full reconcile below.
+    if (
+      previous &&
+      update &&
+      update.index === messages.length - 1 &&
+      update.previousMessage === previous.tailSource &&
+      update.previousMessage.pending &&
+      update.message.pending &&
+      update.previousMessage.id === update.message.id &&
+      update.previousMessage.role === update.message.role &&
+      update.previousMessage.hidden === update.message.hidden &&
+      update.previousMessage.branchGroupId === update.message.branchGroupId &&
+      previous.tailItem?.message.id === update.message.id &&
+      previous.repository.headId === update.message.id
+    ) {
+      const cachedMessage = cacheRef.current.get(update.message)
+      const runtimeMessage = cachedMessage ?? toRuntimeMessage(update.message)
+
+      if (!cachedMessage) {
+        cacheRef.current.set(update.message, runtimeMessage)
+      }
+
+      const tailItem = { message: runtimeMessage, parentId: previous.tailItem.parentId }
+
+      const repository: RuntimeMessageRepository = {
+        headId: update.message.id,
+        messages: previous.repository.messages,
+        [runtimeMessageRepositoryDelta]: tailItem
+      }
+
+      repositoryCacheRef.current = { repository, source: messages, tailItem, tailSource: nextTail }
+
+      return repository
+    }
+
     const items: { message: ThreadMessage; parentId: string | null }[] = []
     const branchParentByGroup = new Map<string, string | null>()
     let visibleParentId: string | null = null
@@ -46,6 +105,13 @@ export function useRuntimeMessageRepository(messages: ChatMessage[]): ExportedMe
       }
     }
 
-    return ExportedMessageRepository.fromBranchableArray(items, { headId })
+    // toRuntimeMessage already returns normalized ThreadMessage objects. Keep
+    // those cached references intact instead of normalizing the whole history a
+    // second time on every streamed delta.
+    const repository: RuntimeMessageRepository = { headId, messages: items }
+    const tailItem = nextTail && items.at(-1)?.message.id === nextTail.id ? (items.at(-1) ?? null) : null
+    repositoryCacheRef.current = { repository, source: messages, tailItem, tailSource: nextTail }
+
+    return repository
   }, [messages])
 }