diff --git a/apps/desktop/src/app/chat/index.tsx b/apps/desktop/src/app/chat/index.tsx index 7a0c0a2674c30..1bb98203b5dff 100644 --- a/apps/desktop/src/app/chat/index.tsx +++ b/apps/desktop/src/app/chat/index.tsx @@ -8,6 +8,7 @@ import { useLocation } from 'react-router' import type { SubmitTextOptions } from '@/app/session/hooks/use-prompt-actions/utils' import { Thread } from '@/components/assistant-ui/thread' +import { TranscriptWindowProvider } from '@/components/assistant-ui/thread/transcript-window' import { Backdrop } from '@/components/Backdrop' import { COMPOSER_HEART_CONFIG, HeartField } from '@/components/chat/vibe-hearts' import { usePaneVisible } from '@/components/pane-shell/pane-visibility' @@ -63,6 +64,7 @@ import { ScrollToBottomButton } from './scroll-to-bottom-button' import { useSessionView } from './session-view' import { SessionActionsMenu } from './sidebar/session-actions-menu' import { threadLoadingState } from './thread-loading' +import { selectTranscriptWindow } from './transcript-window' interface ChatViewProps extends Omit, 'onSubmit'> { gateway: HermesGateway | null @@ -220,9 +222,34 @@ function ChatRuntimeBoundary({ onThreadMessagesChange, suppressMessages }: ChatRuntimeBoundaryProps) { - const storeMessages = useMessagesWhileVisible(useSessionView().$messages) + const view = useSessionView() + const runtimeId = useStore(view.$runtimeId) + const storeMessages = useMessagesWhileVisible(view.$messages) const messages = suppressMessages ? NO_MESSAGES : storeMessages - const runtimeMessageRepository = useRuntimeMessageRepository(messages) + + const [windowPages, setWindowPages] = useState(1) + const [windowSessionKey, setWindowSessionKey] = useState(runtimeId) + + // Reset the window on session swap during RENDER, so a large expand from the + // previous chat can't leak into the next one's first paint (#55191). + if (windowSessionKey !== runtimeId) { + setWindowSessionKey(runtimeId) + setWindowPages(1) + } + + const { messages: windowedMessages, windowed } = useMemo( + () => selectTranscriptWindow(messages, windowPages), + [messages, windowPages] + ) + + const runtimeMessageRepository = useRuntimeMessageRepository(windowedMessages) + + const expandWindow = useCallback(() => setWindowPages(pages => pages + 1), []) + + const transcriptWindow = useMemo( + () => ({ olderAvailable: windowed, expandWindow }), + [expandWindow, windowed] + ) const runtime = useIncrementalExternalStoreRuntime({ messageRepository: runtimeMessageRepository, @@ -237,7 +264,11 @@ function ChatRuntimeBoundary({ onReload }) - return {children} + return ( + + {children} + + ) } // Memoized: the tile caller (session-tile.tsx) and the contrib surface re-render 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 } +} diff --git a/apps/desktop/src/components/assistant-ui/thread/list.test.ts b/apps/desktop/src/components/assistant-ui/thread/list.test.ts index c0cd1da5588a4..a44314d154ffc 100644 --- a/apps/desktop/src/components/assistant-ui/thread/list.test.ts +++ b/apps/desktop/src/components/assistant-ui/thread/list.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from 'vitest' +import { messageRenderWeight, RENDER_WEIGHT_CHARS } from '@/lib/render-weight' + import { buildGroups, firstVisibleGroupIndex, @@ -7,8 +9,6 @@ import { LIVE_TAIL_PARTS, liveTailStart, type MessageGroup, - messageRenderWeight, - RENDER_WEIGHT_CHARS, resolveThreadScrollTarget } from './list' diff --git a/apps/desktop/src/components/assistant-ui/thread/list.tsx b/apps/desktop/src/components/assistant-ui/thread/list.tsx index d8fd0f0a5c665..304b3eb3ec017 100644 --- a/apps/desktop/src/components/assistant-ui/thread/list.tsx +++ b/apps/desktop/src/components/assistant-ui/thread/list.tsx @@ -16,6 +16,7 @@ import { import { type GetTargetScrollTop, useStickToBottom } from 'use-stick-to-bottom' import { useI18n } from '@/i18n' +import { messageRenderWeight } from '@/lib/render-weight' import { cn } from '@/lib/utils' import { onScrollToBottomRequest, @@ -28,6 +29,8 @@ import { isSecondaryWindow } from '@/store/windows' import { MessageRenderBoundary } from '../message-render-boundary' +import { resolveShowEarlierAction, useTranscriptWindow } from './transcript-window' + type ThreadMessageComponents = ComponentProps['components'] export type MessageGroup = { id: string; weight: number } & ( @@ -47,8 +50,6 @@ export type MessageGroup = { id: string; weight: number } & ( // a virtualizer — pure rendering, never touches scrollTop, so it can't fight // use-stick-to-bottom (the single scroll owner). const RENDER_BUDGET = 300 -export const RENDER_WEIGHT_CHARS = 512 -const MAX_MEASURED_MESSAGE_CHARS = RENDER_BUDGET * RENDER_WEIGHT_CHARS // On session switch, paint a small budget first (enough for the bottom turn(s) // the user actually sees after scroll-to-bottom), then bump to the full budget // in a requestAnimationFrame — defers the heavy markdown+syntax-highlight render @@ -76,70 +77,6 @@ export const resolveThreadScrollTarget: GetTargetScrollTop = (targetScrollTop, { return remaining >= 0 && remaining <= SCROLL_TARGET_EPSILON_PX ? currentScrollTop : targetScrollTop } -const contentWeightCache = new WeakMap() -const NON_RENDERED_CONTENT_FIELDS = new Set(['id', 'role', 'toolCallId', 'toolName', 'type']) - -/** - * Estimate the synchronous renderer cost of one assistant-ui message. - * - * The traversal is capped once a single message has enough text to consume a - * complete render page. Going further cannot affect which whole turn crosses - * the budget, and avoiding an unbounded walk matters for deeply nested tool - * payloads. A WeakMap keeps settled history O(message count) on later store - * updates; assistant-ui publishes a new content array when a streaming message - * changes, so the live tail still receives a fresh weight. - */ -export function messageRenderWeight(content: unknown): number { - if (!Array.isArray(content)) { - return 1 - } - - const cached = contentWeightCache.get(content) - - if (cached !== undefined) { - return cached - } - - const seen = new WeakSet() - const pending: unknown[] = [...content] - let characters = 0 - - while (pending.length > 0 && characters < MAX_MEASURED_MESSAGE_CHARS) { - const value = pending.pop() - - if (typeof value === 'string') { - characters += Math.min(value.length, MAX_MEASURED_MESSAGE_CHARS - characters) - - continue - } - - if (!value || typeof value !== 'object' || seen.has(value)) { - continue - } - - seen.add(value) - - if (Array.isArray(value)) { - for (const nested of value) { - pending.push(nested) - } - - continue - } - - for (const [key, nested] of Object.entries(value)) { - if (!NON_RENDERED_CONTENT_FIELDS.has(key)) { - pending.push(nested) - } - } - } - - const weight = Math.max(1, content.length) + Math.ceil(characters / RENDER_WEIGHT_CHARS) - contentWeightCache.set(content, weight) - - return weight -} - interface ThreadMessageListProps { clampToComposer: boolean components: ThreadMessageComponents @@ -315,6 +252,8 @@ const ThreadMessageListInner: FC = ({ targetScrollTop: resolveThreadScrollTarget }) + const { olderAvailable, expandWindow } = useTranscriptWindow() + const [renderBudget, setRenderBudget] = useState(FIRST_PAINT_BUDGET) // Cut the budget during RENDER, not in the post-commit layout effect. An @@ -540,11 +479,26 @@ const ThreadMessageListInner: FC = ({ // Prepend an older page while preserving the on-screen position. The user is // scrolled up (reading history) so the stick-to-bottom lock is escaped and - // won't fight this manual restore. + // won't fight this manual restore. Spend the already-materialized DOM page + // first; only when that is exhausted pull more messages out of the session + // store (#55191). const showEarlier = useCallback(() => { + const action = resolveShowEarlierAction(hiddenCount, olderAvailable) + + if (!action) { + return + } + anchorBeforePrepend() - setRenderBudget(budget => budget + RENDER_BUDGET) - }, [anchorBeforePrepend]) + + if (action === 'dom') { + setRenderBudget(budget => budget + RENDER_BUDGET) + + return + } + + expandWindow() + }, [anchorBeforePrepend, expandWindow, hiddenCount, olderAvailable]) useLayoutEffect(() => { const el = scrollRef.current @@ -553,7 +507,8 @@ const ThreadMessageListInner: FC = ({ el.scrollTop = el.scrollHeight - restoreFromBottomRef.current restoreFromBottomRef.current = null } - }, [scrollRef, renderBudget]) + // renderBudget covers DOM pages; groups.length covers store-window expands. + }, [scrollRef, renderBudget, groups.length]) // The row array is memoized on the inputs the rows actually read. This // component re-renders on every isAtBottom flip — and use-stick-to-bottom @@ -647,7 +602,7 @@ const ThreadMessageListInner: FC = ({ data-slot="aui_thread-content" ref={contentRef as React.RefCallback} > - {hiddenCount > 0 && ( + {(hiddenCount > 0 || olderAvailable) && (