fix(desktop): bound the transcript reaching assistant-ui by render cost (#55191)

An oversized session rebuilt an unbounded runtime repository on every store
update and exhausted the renderer's V8 heap, crash-looping the window. The DOM
budget in thread/list.tsx bounds what PAINTS, but every message was still
normalized into the repository first, so a session only had to be heavy — not
visible — to kill the renderer.

selectTranscriptWindow keeps the tail that fits one render-weight page. Weight,
not message count: measured against a real 1,175-session store, a 400-message
cap disengages on 37 sessions that are heavy but short (one is 133 messages /
1.05MB) while firing on 92 long-but-light sessions that were never at risk.

The cut aligns off branch-group boundaries. useRuntimeMessageRepository records
a group's fork point the first time it sees the group, so a window starting
mid-group would re-parent the surviving branches to whatever happened to
precede them.

Co-authored-by: HexLab <8422520+HexLab98@users.noreply.github.com>

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

@@ -51,4 +51,32 @@ describe('useRuntimeMessageRepository', () => {
 
     expect(feedToRepository(result.current).map(item => item.id)).toEqual(['user-1', 'assistant-stream-1', 'user-2'])
   })
+
+  it('anchors a branch group to its fork point, and a windowed cut keeps it', () => {
+    // Branch groups record their fork parent the first time they are seen. A
+    // window that started mid-group would anchor the survivors to whatever
+    // preceded them instead — selectTranscriptWindow aligns the cut so the
+    // whole group arrives together (#55191).
+    const branch = (id: string): ChatMessage => ({
+      ...text(id, 'assistant', 'branch'),
+      branchGroupId: 'group-1'
+    })
+
+    const messages = [text('user-1', 'user', 'hi'), branch('a-1'), branch('a-2'), text('user-2', 'user', 'more')]
+
+    const { result } = renderHook(() => useRuntimeMessageRepository(messages))
+
+    const parents = new Map(result.current.messages.map(item => [item.message.id, item.parentId]))
+
+    expect(parents.get('a-1')).toBe('user-1')
+    expect(parents.get('a-2')).toBe('user-1')
+
+    // The same group fed as a window that begins AT the group start keeps the
+    // fork intact (parent becomes null: the group is now the transcript root).
+    const { result: windowed } = renderHook(() => useRuntimeMessageRepository(messages.slice(1)))
+
+    const windowedParents = new Map(windowed.current.messages.map(item => [item.message.id, item.parentId]))
+
+    expect(windowedParents.get('a-1')).toBe(windowedParents.get('a-2'))
+  })
 })

apps/desktop/src/app/chat/transcript-window.test.ts

@@ -0,0 +1,144 @@
+import { describe, expect, it } from 'vitest'
+
+import type { ChatMessage } from '@/lib/chat-messages'
+import { RENDER_WEIGHT_CHARS } from '@/lib/render-weight'
+
+import {
+  alignToBranchGroup,
+  selectTranscriptWindow,
+  TRANSCRIPT_WINDOW_BUDGET,
+  TRANSCRIPT_WINDOW_MIN_MESSAGES
+} from './transcript-window'
+
+const message = (id: string, chars: number, branchGroupId?: string): ChatMessage => ({
+  id,
+  parts: [{ type: 'text', text: 'x'.repeat(chars) }],
+  role: id.startsWith('u') ? 'user' : 'assistant',
+  ...(branchGroupId ? { branchGroupId } : {})
+})
+
+/** Messages of `chars` each, newest last. */
+const transcript = (count: number, chars: number): ChatMessage[] =>
+  Array.from({ length: count }, (_, i) => message(`m-${i}`, chars))
+
+describe('selectTranscriptWindow', () => {
+  it('does not window a transcript that fits the budget', () => {
+    const messages = transcript(50, 100)
+
+    const window = selectTranscriptWindow(messages)
+
+    expect(window.windowed).toBe(false)
+    // Reference identity preserved — a fresh array would re-render the runtime.
+    expect(window.messages).toBe(messages)
+  })
+
+  it('windows a HEAVY-but-SHORT transcript that a message-count cap would miss', () => {
+    // 40 messages, each a big tool result. Well under any sane count cap, but
+    // this is the shape that exhausts the renderer heap (#55191).
+    const messages = transcript(40, RENDER_WEIGHT_CHARS * 400)
+
+    const window = selectTranscriptWindow(messages)
+
+    expect(window.windowed).toBe(true)
+    expect(window.messages.length).toBeLessThan(messages.length)
+    expect(window.messages.at(-1)).toBe(messages.at(-1))
+  })
+
+  it('keeps far MORE messages when they are light than when they are heavy', () => {
+    // The contract is weight, not count: a message-count cap would treat these
+    // two identically. 500 tiny messages are cheaper than 500 tool results, so
+    // many more of them survive the same budget.
+    const light = selectTranscriptWindow(transcript(500, 20))
+    const heavy = selectTranscriptWindow(transcript(500, RENDER_WEIGHT_CHARS * 40))
+
+    expect(light.messages.length).toBeGreaterThan(heavy.messages.length * 10)
+  })
+
+  it('leaves a long transcript whole when the whole thing is cheap', () => {
+    const messages = transcript(600, 20)
+
+    const window = selectTranscriptWindow(messages)
+
+    expect(window.windowed).toBe(false)
+    expect(window.messages).toBe(messages)
+  })
+
+  it('keeps a floor of messages when single turns are enormous', () => {
+    const messages = transcript(80, RENDER_WEIGHT_CHARS * TRANSCRIPT_WINDOW_BUDGET)
+
+    const window = selectTranscriptWindow(messages)
+
+    expect(window.messages.length).toBeGreaterThanOrEqual(TRANSCRIPT_WINDOW_MIN_MESSAGES)
+  })
+
+  it('grows by one budget page per expand and eventually covers everything', () => {
+    const messages = transcript(400, RENDER_WEIGHT_CHARS * 40)
+
+    const first = selectTranscriptWindow(messages, 1)
+    const second = selectTranscriptWindow(messages, 2)
+
+    expect(first.windowed).toBe(true)
+    expect(second.messages.length).toBeGreaterThan(first.messages.length)
+
+    let pages = 1
+    let window = selectTranscriptWindow(messages, pages)
+
+    while (window.windowed && pages < 100) {
+      window = selectTranscriptWindow(messages, ++pages)
+    }
+
+    // Paging terminates at the full transcript — never a dead end.
+    expect(window.windowed).toBe(false)
+    expect(window.messages).toHaveLength(messages.length)
+  })
+
+  it('never cuts inside a branch group, so branches keep their fork point', () => {
+    const heavy = RENDER_WEIGHT_CHARS * 200
+
+    // A branch group sits right where a weight-only cut would land.
+    const messages: ChatMessage[] = [
+      ...transcript(20, heavy),
+      message('a-branch-1', heavy, 'group-1'),
+      message('a-branch-2', heavy, 'group-1'),
+      message('a-branch-3', heavy, 'group-1'),
+      ...transcript(20, heavy).map(m => ({ ...m, id: `tail-${m.id}` }))
+    ]
+
+    for (let pages = 1; pages <= 6; pages++) {
+      const kept = selectTranscriptWindow(messages, pages).messages
+      const groupMembers = kept.filter(m => m.branchGroupId === 'group-1')
+
+      // Either the whole group survives or none of it does — never a partial
+      // group, which would re-parent the surviving branches.
+      expect([0, 3]).toContain(groupMembers.length)
+    }
+  })
+
+  it('handles an empty transcript', () => {
+    const messages: ChatMessage[] = []
+
+    expect(selectTranscriptWindow(messages)).toEqual({ messages, windowed: false })
+  })
+})
+
+describe('alignToBranchGroup', () => {
+  const messages = [
+    message('u-1', 10),
+    message('a-1', 10, 'g'),
+    message('a-2', 10, 'g'),
+    message('u-2', 10)
+  ]
+
+  it('widens a cut that lands mid-group back to the group start', () => {
+    expect(alignToBranchGroup(messages, 2)).toBe(1)
+  })
+
+  it('leaves a cut on a non-branch message alone', () => {
+    expect(alignToBranchGroup(messages, 3)).toBe(3)
+  })
+
+  it('clamps out-of-range indices', () => {
+    expect(alignToBranchGroup(messages, -5)).toBe(0)
+    expect(alignToBranchGroup(messages, 99)).toBe(messages.length)
+  })
+})

apps/desktop/src/app/chat/transcript-window.ts

@@ -0,0 +1,111 @@
+import type { ChatMessage } from '@/lib/chat-messages'
+import { messageRenderWeight } from '@/lib/render-weight'
+
+/**
+ * Bound what reaches assistant-ui at all.
+ *
+ * Rendering the full transcript of an oversized session rebuilds an unbounded
+ * runtime repository on every store update and exhausts the renderer's V8 heap
+ * (#55191). The DOM budget in `thread/list.tsx` already bounds what PAINTS, but
+ * every message still gets normalized into the repository first — so a session
+ * only has to be heavy, not visible, to crash the window.
+ *
+ * The window spends the same currency as the DOM budget: render weight, not
+ * message count. A count cap gets this wrong in both directions — measured on a
+ * real 1,175-session store, a 400-message cap would disable itself on 37
+ * sessions that are heavy but short (one was 133 messages / 1.05MB) while
+ * engaging on 92 long-but-light sessions that were never at risk.
+ */
+
+/**
+ * One window page, in render-weight units.
+ *
+ * Four DOM pages (the `RENDER_BUDGET` of 300 in `thread/list.tsx`). "Show
+ * earlier" spends the DOM budget first, so the user pages through the
+ * already-materialized window three times before this asks the store for more
+ * — and the reported crash shape (~231K tokens ≈ 2,260 units) is windowed
+ * rather than handed to the repository whole.
+ */
+export const TRANSCRIPT_WINDOW_BUDGET = 1200
+
+/**
+ * Floor on messages kept regardless of weight. A transcript of enormous turns
+ * must still render the turn the user is having; without this a single
+ * multi-megabyte tool result could window everything after it away.
+ */
+export const TRANSCRIPT_WINDOW_MIN_MESSAGES = 30
+
+export interface TranscriptWindow {
+  /** The tail assistant-ui is allowed to materialize. */
+  messages: ChatMessage[]
+  /** Store holds older messages than this window. */
+  windowed: boolean
+}
+
+/**
+ * Widen a cut backwards so it never lands inside an assistant branch group.
+ *
+ * `useRuntimeMessageRepository` records a group's fork point the first time it
+ * sees the group (`branchParentByGroup`). A cut through the middle of a group
+ * therefore anchors the surviving branches to whatever message happens to
+ * precede them in the window — silently re-parenting a branch. Include the
+ * whole group or none of it.
+ */
+export function alignToBranchGroup(messages: readonly ChatMessage[], start: number): number {
+  if (start <= 0 || start >= messages.length) {
+    return Math.max(0, Math.min(start, messages.length))
+  }
+
+  const group = messages[start].branchGroupId
+
+  if (!group) {
+    return start
+  }
+
+  let aligned = start
+
+  while (aligned > 0 && messages[aligned - 1].branchGroupId === group) {
+    aligned--
+  }
+
+  return aligned
+}
+
+/**
+ * Select the tail of the transcript that fits one window, grown by `pages`.
+ *
+ * Walks newest-first accumulating weight until the budget is met, keeps at
+ * least MIN messages, then aligns the cut off a branch-group boundary.
+ */
+export function selectTranscriptWindow(
+  messages: readonly ChatMessage[],
+  pages = 1
+): TranscriptWindow {
+  const budget = TRANSCRIPT_WINDOW_BUDGET * Math.max(1, Math.floor(pages))
+
+  if (messages.length === 0) {
+    return { messages: messages as ChatMessage[], windowed: false }
+  }
+
+  let start = messages.length
+  let weight = 0
+
+  for (let i = messages.length - 1; i >= 0; i--) {
+    weight += messageRenderWeight(messages[i].parts)
+    start = i
+
+    if (weight >= budget && messages.length - i >= TRANSCRIPT_WINDOW_MIN_MESSAGES) {
+      break
+    }
+  }
+
+  start = alignToBranchGroup(messages, start)
+
+  if (start <= 0) {
+    // Preserve reference identity when the whole transcript fits: handing React
+    // a fresh array of the same messages re-renders the runtime for nothing.
+    return { messages: messages as ChatMessage[], windowed: false }
+  }
+
+  return { messages: messages.slice(start), windowed: true }
+}