Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 34 additions & 3 deletions apps/desktop/src/app/chat/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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<React.ComponentProps<'div'>, 'onSubmit'> {
gateway: HermesGateway | null
Expand Down Expand Up @@ -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<ThreadMessage>({
messageRepository: runtimeMessageRepository,
Expand All @@ -237,7 +264,11 @@ function ChatRuntimeBoundary({
onReload
})

return <AssistantRuntimeProvider runtime={runtime}>{children}</AssistantRuntimeProvider>
return (
<TranscriptWindowProvider value={transcriptWindow}>
<AssistantRuntimeProvider runtime={runtime}>{children}</AssistantRuntimeProvider>
</TranscriptWindowProvider>
)
}

// Memoized: the tile caller (session-tile.tsx) and the contrib surface re-render
Expand Down
28 changes: 28 additions & 0 deletions apps/desktop/src/app/chat/runtime-repository.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'))
})
})
144 changes: 144 additions & 0 deletions apps/desktop/src/app/chat/transcript-window.test.ts
Original file line number Diff line number Diff line change
@@ -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)
})
})
111 changes: 111 additions & 0 deletions apps/desktop/src/app/chat/transcript-window.ts
Original file line number Diff line number Diff line change
@@ -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 }
}
4 changes: 2 additions & 2 deletions apps/desktop/src/components/assistant-ui/thread/list.test.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
import { describe, expect, it } from 'vitest'

import { messageRenderWeight, RENDER_WEIGHT_CHARS } from '@/lib/render-weight'

import {
buildGroups,
firstVisibleGroupIndex,
LIVE_TAIL_MIN_GROUPS,
LIVE_TAIL_PARTS,
liveTailStart,
type MessageGroup,
messageRenderWeight,
RENDER_WEIGHT_CHARS,
resolveThreadScrollTarget
} from './list'

Expand Down
Loading
Loading