diff --git a/apps/desktop/src/app/chat/runtime-repository.test.ts b/apps/desktop/src/app/chat/runtime-repository.test.ts index ae0d4bd469ef2..de7a97d39c575 100644 --- a/apps/desktop/src/app/chat/runtime-repository.test.ts +++ b/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')) + }) }) diff --git a/apps/desktop/src/app/chat/transcript-window.test.ts b/apps/desktop/src/app/chat/transcript-window.test.ts new file mode 100644 index 0000000000000..4726724f68e26 --- /dev/null +++ b/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) + }) +}) diff --git a/apps/desktop/src/app/chat/transcript-window.ts b/apps/desktop/src/app/chat/transcript-window.ts new file mode 100644 index 0000000000000..29d2d1c9d2783 --- /dev/null +++ b/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 } +}