Merge pull request #71835 from NousResearch/bb/stream-history-cost

perf(desktop): make streaming cost independent of transcript length

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

@@ -8,11 +8,13 @@ import multitab from './multitab.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,
   [multitab.name]: multitab,

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

@@ -0,0 +1,18 @@
+// Streaming into an ALREADY-LONG transcript. Same measurement as `stream`, but
+// the history is mounted and allowed to settle before the recorders start, so
+// what it captures is the per-delta cost that scales with transcript length —
+// the regression reported in #69120.
+//
+// Report-only (tier: manual): the number depends on how much history the host
+// can mount, so it is not gated against the committed baseline.
+
+import stream from './stream.mjs'
+
+export default {
+  name: 'stream-history',
+  tier: 'manual',
+  description: 'Streaming cost with a long settled transcript already mounted.',
+  run(cdp, opts = {}) {
+    return stream.run(cdp, { ...opts, historyTurns: Number(opts.historyTurns ?? 200) })
+  }
+}

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

@@ -70,7 +70,7 @@ const COLLECT = `
   })()
 `
 
-function analyze(data, warmupMs) {
+function analyze(data, warmupMs, extra = {}) {
   // Drop warm-up frames (recorder installs before the stream starts).
   const frames = []
   let acc = 0
@@ -102,6 +102,7 @@ function analyze(data, warmupMs) {
       intermut_p95_ms: Math.round(percentile(interMut, 0.95) * 10) / 10
     },
     detail: {
+      ...extra,
       windowS: Math.round(windowS * 10) / 10,
       avgFps: windowS ? Math.round((frames.length / windowS) * 10) / 10 : 0,
       frameHistogram: frameHistogram(frames),
@@ -128,13 +129,36 @@ export default {
     // noise unrelated to render cost).
     const chunk = opts.chunk ?? 'A streamed sentence with **bold**, `code`, and ordinary prose like a normal reply.\n\n'
     const real = Boolean(opts.real)
+    const historyTurns = Number(opts.historyTurns ?? 0)
+    const historySettleMs = Number(opts.historySettleMs ?? 1500)
 
     await cdp.send('Runtime.enable')
+
+    // Mount the settled history BEFORE the recorders start, so the measurement
+    // window contains only streaming work — not the one-off mount cost.
+    if (historyTurns > 0) {
+      if (real) {
+        throw new Error('--historyTurns is only supported by the synthetic stream path')
+      }
+
+      await cdp.eval(`window.__PERF_DRIVE__.loadTranscript(${historyTurns})`)
+      await sleep(historySettleMs)
+
+      const mounted = Number(await cdp.eval('window.__PERF_DRIVE__.snapshotMsgs()'))
+      const expected = historyTurns * 2
+
+      if (mounted !== expected) {
+        throw new Error(`expected ${expected} preloaded history messages, got ${mounted}`)
+      }
+    }
+
     await cdp.eval(RECORDERS)
 
     if (real) {
       // Backend path: fire a real prompt and wait for the stream to appear.
-      const baseCount = await cdp.eval(`document.querySelectorAll(${JSON.stringify(SELECTORS.assistantMessage)}).length`)
+      const baseCount = await cdp.eval(
+        `document.querySelectorAll(${JSON.stringify(SELECTORS.assistantMessage)}).length`
+      )
       await typeIntoComposer(cdp, opts.prompt ?? 'count from 1 to 80, one number per line', { cps: 40 })
       await cdp.eval(`(() => {
         const el = document.querySelector(${JSON.stringify(SELECTORS.composer)})
@@ -186,6 +210,6 @@ export default {
       await cdp.eval('window.__PERF_DRIVE__.reset()')
     }
 
-    return analyze(data, real ? 0 : 500)
+    return analyze(data, real ? 0 : 500, { historyTurns })
   }
 }

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

@@ -1,14 +1,27 @@
-import { ExportedMessageRepository, type ThreadMessage } from '@assistant-ui/react'
+import { fromThreadMessageLike, getAutoStatus } from '@assistant-ui/core/internal'
+import type { ExportedMessageRepository, ThreadMessage } from '@assistant-ui/react'
 import { useMemo, useRef } from 'react'
 
 import type { ChatMessage } from '@/lib/chat-messages'
 import { coalesceToolOnlyAssistants, createToolMergeCache, toRuntimeMessage } from '@/lib/chat-runtime'
 
+// The exact fallback status ExportedMessageRepository.fromBranchableArray uses.
+// Normalization happens HERE, once per message, so the cached record below is
+// already the final ThreadMessage the runtime consumes.
+const FALLBACK_STATUS = getAutoStatus(false, false, false, false, undefined)
+
 /**
  * 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.
+ *
+ * The cache stores NORMALIZED messages. `fromBranchableArray` maps the whole
+ * array through `fromThreadMessageLike` on every call, so building the export
+ * with it threw away the cache's reference identity once per streamed delta —
+ * re-normalizing the entire settled transcript ~30x/s. Normalizing inside the
+ * cache miss keeps identity stable for settled turns, which is what lets the
+ * runtime reconcile detect that only the tail moved.
  */
 export function useRuntimeMessageRepository(messages: ChatMessage[]): ExportedMessageRepository {
   const cacheRef = useRef(new WeakMap<ChatMessage, ThreadMessage>())
@@ -32,7 +45,9 @@ export function useRuntimeMessageRepository(messages: ChatMessage[]): ExportedMe
       }
 
       const cachedMessage = cacheRef.current.get(message)
-      const runtimeMessage = cachedMessage ?? toRuntimeMessage(message)
+
+      const runtimeMessage =
+        cachedMessage ?? fromThreadMessageLike(toRuntimeMessage(message), message.id, FALLBACK_STATUS)
 
       if (!cachedMessage) {
         cacheRef.current.set(message, runtimeMessage)
@@ -46,6 +61,6 @@ export function useRuntimeMessageRepository(messages: ChatMessage[]): ExportedMe
       }
     }
 
-    return ExportedMessageRepository.fromBranchableArray(items, { headId })
+    return { headId, messages: items }
   }, [messages])
 }

apps/desktop/src/app/session/hooks/use-message-stream/index.ts

@@ -251,12 +251,20 @@ export function useMessageStream({
       flushQueuedDeltas()
     }
 
-    if (sinceLast >= STREAM_DELTA_FLUSH_MS && typeof window.requestAnimationFrame === 'function') {
-      flushHandleRef.current = window.requestAnimationFrame(runFlush)
-
-      return
-    }
-
+    // Always a timer, never requestAnimationFrame. Chromium pauses rAF for a
+    // renderer it considers hidden, and "hidden" is not something this code can
+    // verify: `backgroundThrottling: false` plus the process-level switches in
+    // electron/main.ts cover the blurred and occluded cases, but they don't
+    // cover a minimized window, a fully off-screen one, or a renderer the
+    // compositor has otherwise parked. In those states an rAF-gated flush never
+    // runs, so a finished answer sits in this queue until some later input or
+    // focus event happens to wake a frame — the reply looks stalled, then
+    // arrives all at once on refocus.
+    //
+    // A timer keeps the same coalescing cadence (that's what the floor above is
+    // for) while guaranteeing delivery without user interaction. Timers are
+    // clamped in background renderers rather than suspended, and
+    // disable-background-timer-throttling already opts out of that clamp.
     flushHandleRef.current = window.setTimeout(runFlush, Math.max(0, STREAM_DELTA_FLUSH_MS - sinceLast))
   }, [flushQueuedDeltas])
 
@@ -277,11 +285,7 @@ export function useMessageStream({
   useEffect(
     () => () => {
       if (flushHandleRef.current !== null && typeof window !== 'undefined') {
-        if (typeof window.cancelAnimationFrame === 'function') {
-          window.cancelAnimationFrame(flushHandleRef.current)
-        } else {
-          window.clearTimeout(flushHandleRef.current)
-        }
+        window.clearTimeout(flushHandleRef.current)
       }
 
       flushHandleRef.current = null

apps/desktop/src/app/session/hooks/use-message-stream/stream-flush.test.tsx

@@ -0,0 +1,93 @@
+import { QueryClient } from '@tanstack/react-query'
+import { act, cleanup, render } from '@testing-library/react'
+import { useEffect, useRef } from 'react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import type { ClientSessionState } from '@/app/types'
+import { createClientSessionState } from '@/lib/chat-runtime'
+import type { RpcEvent } from '@/types/hermes'
+
+import { STREAM_DELTA_FLUSH_MS } from './utils'
+
+import { useMessageStream } from './index'
+
+const SID = 'stream-session'
+
+let handleEvent: ((event: RpcEvent) => void) | null = null
+let states: Map<string, ClientSessionState>
+
+function Harness() {
+  const activeSessionIdRef = useRef<string | null>(SID)
+  const sessionStateByRuntimeIdRef = useRef(new Map<string, ClientSessionState>())
+  const queryClientRef = useRef(new QueryClient())
+
+  const stream = useMessageStream({
+    activeSessionIdRef,
+    hydrateFromStoredSession: vi.fn(async () => undefined),
+    queryClient: queryClientRef.current,
+    refreshHermesConfig: vi.fn(async () => undefined),
+    refreshSessions: vi.fn(async () => undefined),
+    sessionStateByRuntimeIdRef,
+    updateSessionState: (sessionId, updater) => {
+      const current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState()
+      const next = updater(current)
+      sessionStateByRuntimeIdRef.current.set(sessionId, next)
+
+      return next
+    }
+  })
+
+  useEffect(() => {
+    handleEvent = stream.handleGatewayEvent
+    states = sessionStateByRuntimeIdRef.current
+  }, [stream.handleGatewayEvent])
+
+  return null
+}
+
+describe('stream delta delivery', () => {
+  beforeEach(() => {
+    handleEvent = null
+    states = new Map()
+  })
+
+  afterEach(() => {
+    cleanup()
+    vi.useRealTimers()
+    vi.restoreAllMocks()
+  })
+
+  it('delivers a delta while animation frames are parked', async () => {
+    // An Electron renderer the compositor has parked: rAF accepts the callback
+    // but never runs it. A frame-gated flush strands the answer here until some
+    // later focus/input event wakes a frame.
+    const rafSpy = vi.spyOn(window, 'requestAnimationFrame').mockImplementation(() => 1)
+    vi.useFakeTimers()
+
+    render(<Harness />)
+    await act(async () => {
+      await Promise.resolve()
+    })
+
+    // Reach the frame-gated branch: it is only taken once the coalescing floor
+    // has already elapsed since the previous flush. Send one delta, let it
+    // flush, then idle past the floor so the NEXT delta schedules immediately.
+    act(() => handleEvent?.({ payload: { text: 'first ' }, session_id: SID, type: 'message.delta' }))
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(STREAM_DELTA_FLUSH_MS)
+    })
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(STREAM_DELTA_FLUSH_MS * 2)
+    })
+
+    act(() => handleEvent?.({ payload: { text: 'and the rest' }, session_id: SID, type: 'message.delta' }))
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(STREAM_DELTA_FLUSH_MS)
+    })
+
+    expect(states.get(SID)?.messages.at(-1)?.parts).toEqual([{ type: 'text', text: 'first and the rest' }])
+    // The flush must not have depended on a frame at all.
+    expect(rafSpy).not.toHaveBeenCalled()
+  })
+})

apps/desktop/src/lib/incremental-external-store-runtime.test.ts

@@ -0,0 +1,144 @@
+import { fromThreadMessageLike, getAutoStatus, MessageRepository } from '@assistant-ui/core/internal'
+import type { ExportedMessageRepository, ThreadMessage } from '@assistant-ui/react'
+import { describe, expect, it, vi } from 'vitest'
+
+import { syncRepositoryIncrementally } from './incremental-external-store-runtime'
+
+const STATUS = getAutoStatus(false, false, false, false, undefined)
+
+function message(id: string, text: string): ThreadMessage {
+  return fromThreadMessageLike({ role: 'assistant', content: [{ type: 'text', text }] }, id, STATUS)
+}
+
+/** A real MessageRepository behind the same shape syncRepositoryIncrementally drives. */
+function runtimeWith(items: { message: ThreadMessage; parentId: string | null }[]) {
+  const repository = new MessageRepository()
+
+  for (const { message: item, parentId } of items) {
+    repository.addOrUpdateMessage(parentId, item)
+  }
+
+  if (items.length > 0) {
+    repository.resetHead(items.at(-1)?.message.id ?? null)
+  }
+
+  return { repository } as unknown as Parameters<typeof syncRepositoryIncrementally>[0]
+}
+
+function chain(messages: ThreadMessage[]) {
+  return messages.map((item, index) => ({
+    message: item,
+    parentId: index === 0 ? null : messages[index - 1].id
+  }))
+}
+
+function exported(items: { message: ThreadMessage; parentId: string | null }[]): ExportedMessageRepository {
+  return { headId: items.at(-1)?.message.id ?? null, messages: items }
+}
+
+describe('syncRepositoryIncrementally', () => {
+  it('writes only the changed tail instead of the whole transcript', () => {
+    const settled = Array.from({ length: 200 }, (_, index) => message(`m-${index}`, `body ${index}`))
+    const items = chain(settled)
+    const runtime = runtimeWith(items)
+    const repository = (runtime as unknown as { repository: MessageRepository }).repository
+
+    const addOrUpdate = vi.spyOn(repository, 'addOrUpdateMessage')
+    const resetHead = vi.spyOn(repository, 'resetHead')
+
+    // One streamed delta: the tail grows, every settled message keeps identity.
+    const nextTail = message('m-199', 'body 199 + delta')
+    const nextItems = [...items.slice(0, -1), { message: nextTail, parentId: 'm-198' }]
+
+    const result = syncRepositoryIncrementally(runtime, exported(nextItems))
+
+    expect(addOrUpdate).toHaveBeenCalledTimes(1)
+    expect(addOrUpdate).toHaveBeenCalledWith('m-198', nextTail)
+    // The head did not move, so the descendant-pruning reset is skipped.
+    expect(resetHead).not.toHaveBeenCalled()
+    expect(result).toHaveLength(200)
+    expect(result.at(-1)).toBe(nextTail)
+  })
+
+  it('does nothing at all when the transcript is unchanged', () => {
+    const items = chain([message('a', 'one'), message('b', 'two')])
+    const runtime = runtimeWith(items)
+    const repository = (runtime as unknown as { repository: MessageRepository }).repository
+
+    const addOrUpdate = vi.spyOn(repository, 'addOrUpdateMessage')
+    const deleteMessage = vi.spyOn(repository, 'deleteMessage')
+
+    syncRepositoryIncrementally(runtime, exported(items))
+
+    expect(addOrUpdate).not.toHaveBeenCalled()
+    expect(deleteMessage).not.toHaveBeenCalled()
+  })
+
+  it('appends a new message through the full path', () => {
+    const first = message('a', 'one')
+    const items = chain([first])
+    const runtime = runtimeWith(items)
+
+    const second = message('b', 'two')
+    const result = syncRepositoryIncrementally(runtime, exported(chain([first, second])))
+
+    expect(result.map(item => item.id)).toEqual(['a', 'b'])
+  })
+
+  it('honours an authoritative deletion', () => {
+    const a = message('a', 'one')
+    const b = message('b', 'two')
+    const c = message('c', 'three')
+    const runtime = runtimeWith(chain([a, b, c]))
+
+    const result = syncRepositoryIncrementally(runtime, exported(chain([a, b])))
+
+    expect(result.map(item => item.id)).toEqual(['a', 'b'])
+  })
+
+  it('rebuilds cleanly when a disjoint transcript is swapped in', () => {
+    const runtime = runtimeWith(chain([message('old-1', 'one'), message('old-2', 'two')]))
+
+    const next = chain([message('new-1', 'alpha'), message('new-2', 'beta')])
+    const result = syncRepositoryIncrementally(runtime, exported(next))
+
+    expect(result.map(item => item.id)).toEqual(['new-1', 'new-2'])
+  })
+
+  it('re-parents a message when its branch parent changes', () => {
+    const root = message('root', 'root')
+    const a = message('a', 'a')
+    const b = message('b', 'b')
+
+    const runtime = runtimeWith([
+      { message: root, parentId: null },
+      { message: a, parentId: 'root' },
+      { message: b, parentId: 'a' }
+    ])
+
+    // Same ids and same message objects, but `b` moves onto a sibling branch.
+    const result = syncRepositoryIncrementally(runtime, {
+      headId: 'b',
+      messages: [
+        { message: root, parentId: null },
+        { message: a, parentId: 'root' },
+        { message: b, parentId: 'root' }
+      ]
+    })
+
+    expect(result.map(item => item.id)).toEqual(['root', 'b'])
+  })
+
+  it('moves the head when an explicit headId rewinds the branch', () => {
+    const a = message('a', 'one')
+    const b = message('b', 'two')
+    const runtime = runtimeWith(chain([a, b]))
+
+    const result = syncRepositoryIncrementally(runtime, {
+      headId: 'a',
+      messages: chain([a, b])
+    })
+
+    expect(result.map(item => item.id)).toEqual(['a'])
+  })
+})

apps/desktop/src/lib/incremental-external-store-runtime.ts

@@ -35,25 +35,81 @@ const shallowEqual = (a: object, b: object): boolean => {
 
 const getThreadListAdapter = (store: ExternalStoreAdapter) => store.adapters?.threadList ?? {}
 
-function syncRepositoryIncrementally(
+/**
+ * Write only the items whose (message, parentId) pair actually moved.
+ *
+ * `useRuntimeMessageRepository` caches normalized ThreadMessages by source
+ * identity, so a settled turn keeps the SAME object across renders. That makes
+ * an identity check a sound "did this change?" test: during streaming exactly
+ * one item — the growing tail — differs, and the other N-1 writes were pure
+ * overhead that grew with transcript length.
+ *
+ * Returns false when the export is stale (an id in `existing` is gone, or an
+ * incoming message has no repository entry yet), so the caller falls back to
+ * the full rebuild rather than guessing.
+ */
+function applyChangedMessages(
+  repository: ExternalStoreThreadRuntimeCore['repository'],
+  existing: readonly { message: ThreadMessage; parentId: string | null }[],
+  incoming: readonly { message: ThreadMessage; parentId: string | null }[]
+): boolean {
+  if (existing.length !== incoming.length) {
+    return false
+  }
+
+  const existingById = new Map(existing.map(item => [item.message.id, item]))
+
+  for (const item of incoming) {
+    const current = existingById.get(item.message.id)
+
+    if (!current) {
+      return false
+    }
+
+    // Reference identity, not deep equality: the conversion cache guarantees a
+    // stable object for an unchanged turn, and a changed turn is a new object.
+    if (current.message !== item.message || current.parentId !== item.parentId) {
+      repository.addOrUpdateMessage(item.parentId, item.message)
+    }
+  }
+
+  return true
+}
+
+export function syncRepositoryIncrementally(
   runtime: ExternalStoreThreadRuntimeCore,
   messageRepository: NonNullable<ExternalStoreAdapter['messageRepository']>
 ): readonly ThreadMessage[] {
   const repository = (runtime as unknown as { repository: ExternalStoreThreadRuntimeCore['repository'] }).repository
-  const incomingIds = new Set(messageRepository.messages.map(({ message }) => message.id))
+  const incoming = messageRepository.messages
   const existing = repository.export().messages
+  const headId = messageRepository.headId ?? incoming.at(-1)?.message.id ?? null
 
   // A thread switch swaps in a fully-DISJOINT transcript (no id carries over).
   // Reconciling two unrelated trees in place — grafting the new chain onto the
   // old one, then pruning — can strand a stale head/branch, so there's nothing
   // to preserve: clear the tree first (leaves→root), then rebuild clean.
-  if (existing.length > 0 && !existing.some(({ message }) => incomingIds.has(message.id))) {
+  const incomingIds = new Set(incoming.map(({ message }) => message.id))
+  const disjoint = existing.length > 0 && !existing.some(({ message }) => incomingIds.has(message.id))
+
+  // Steady-state streaming: same message set, one item changed. Skip the
+  // whole-transcript rewrite, the prune scan, and the second export. resetHead
+  // deletes the head's descendants, so it only runs when the head really moved.
+  if (!disjoint && applyChangedMessages(repository, existing, incoming)) {
+    if (repository.headId !== headId) {
+      repository.resetHead(headId)
+    }
+
+    return repository.getMessages()
+  }
+
+  if (disjoint) {
     for (const { message } of [...existing].reverse()) {
       repository.deleteMessage(message.id)
     }
   }
 
-  for (const { message, parentId } of messageRepository.messages) {
+  for (const { message, parentId } of incoming) {
     repository.addOrUpdateMessage(parentId, message)
   }
 
@@ -63,8 +119,6 @@ function syncRepositoryIncrementally(
     }
   }
 
-  const headId = messageRepository.headId ?? messageRepository.messages.at(-1)?.message.id ?? null
-
   repository.resetHead(headId)
 
   return repository.getMessages()