From 64ab17182a9052184b167b7a7100a2da7cb3134b Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Sat, 16 May 2026 21:17:36 -0500 Subject: [PATCH] feat(desktop): virtualize chat thread + sidebar via TanStack Virtual MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces `use-stick-to-bottom` and per-row session rendering with `@tanstack/react-virtual`, matching what Cursor uses. Chat thread (`thread-virtualizer.tsx`): - Natural-flow virtualization (padding spacers, not absolute items) so `position: sticky` on the human bubble still resolves cleanly against the scroller. - Custom at-bottom anchor: pins when armed, disarms on user-driven upward scroll, re-arms at bottom, jumps on session switch + `thread.runStart`. - Loading indicator and `--thread-last-message-clearance` move to a real `[data-slot=aui_composer-clearance]` node; drops the brittle `:nth-last-child(1 of …)` rule that can't fire reliably under virtualization. Sidebar (`virtual-session-list.tsx`): - Flat agents list virtualizes at >=25 rows; pinned and workspace-grouped paths stay direct-render. - `SortableContext` keeps all IDs; only the window mounts; dnd-kit's `setNodeRef` is merged with `virtualizer.measureElement` so rows participate in both DnD hit-testing and TanStack measurement. Drops `use-stick-to-bottom`. Streaming test gets a global `offsetWidth/offsetHeight` stub so the virtualizer's viewport sizing works in jsdom; the scroll-up-doesn't-pull-back invariant still passes. --- apps/desktop/package.json | 2 +- apps/desktop/src/app/chat/sidebar/index.tsx | 24 +- .../app/chat/sidebar/virtual-session-list.tsx | 149 +++++++++ .../assistant-ui/streaming.test.tsx | 65 ++++ .../assistant-ui/thread-virtualizer.tsx | 306 +++++++++++++++++ .../src/components/assistant-ui/thread.tsx | 311 ++---------------- apps/desktop/src/styles.css | 17 +- package-lock.json | 36 +- 8 files changed, 592 insertions(+), 318 deletions(-) create mode 100644 apps/desktop/src/app/chat/sidebar/virtual-session-list.tsx create mode 100644 apps/desktop/src/components/assistant-ui/thread-virtualizer.tsx diff --git a/apps/desktop/package.json b/apps/desktop/package.json index 6f6660744ab..634356f55f0 100644 --- a/apps/desktop/package.json +++ b/apps/desktop/package.json @@ -58,6 +58,7 @@ "@tailwindcss/typography": "^0.5.19", "@tailwindcss/vite": "^4.2.4", "@tanstack/react-query": "^5.100.6", + "@tanstack/react-virtual": "^3.13.24", "@vscode/codicons": "^0.0.45", "@xterm/addon-fit": "^0.11.0", "@xterm/addon-unicode11": "^0.9.0", @@ -88,7 +89,6 @@ "unicode-animations": "^1.0.3", "unified": "^11.0.5", "unist-util-visit-parents": "^6.0.2", - "use-stick-to-bottom": "^1.1.4", "vfile": "^6.0.3", "web-haptics": "^0.0.6" }, diff --git a/apps/desktop/src/app/chat/sidebar/index.tsx b/apps/desktop/src/app/chat/sidebar/index.tsx index 0888859ed86..05cc5b41dc4 100644 --- a/apps/desktop/src/app/chat/sidebar/index.tsx +++ b/apps/desktop/src/app/chat/sidebar/index.tsx @@ -62,6 +62,9 @@ import { SidebarPanelLabel } from '../../shell/sidebar-label' import type { SidebarNavItem } from '../../types' import { SidebarSessionRow } from './session-row' +import { VirtualSessionList } from './virtual-session-list' + +const VIRTUALIZE_THRESHOLD = 25 const SIDEBAR_NAV: SidebarNavItem[] = [ { id: 'new-session', label: 'New agent', icon: props => , action: 'new-session' }, @@ -539,6 +542,8 @@ function SidebarSessionsSection({ renderRows(items) ) + const flatVirtualized = !showEmptyState && !groups?.length && sessions.length >= VIRTUALIZE_THRESHOLD + let inner: React.ReactNode if (showEmptyState) { @@ -559,6 +564,19 @@ function SidebarSessionsSection({ ) : ( groupNodes ) + } else if (flatVirtualized) { + inner = ( + + ) } else { inner = renderSessionList(sessions) } @@ -572,11 +590,15 @@ function SidebarSessionsSection({ inner ) + // The virtualizer owns its own scroller, so suppress the wrapper's overflow + // to avoid a double scroll container. + const resolvedContentClassName = cn(contentClassName, flatVirtualized && 'overflow-y-visible') + return ( {open && ( - + {body} {footer} diff --git a/apps/desktop/src/app/chat/sidebar/virtual-session-list.tsx b/apps/desktop/src/app/chat/sidebar/virtual-session-list.tsx new file mode 100644 index 00000000000..7613c621736 --- /dev/null +++ b/apps/desktop/src/app/chat/sidebar/virtual-session-list.tsx @@ -0,0 +1,149 @@ +import { SortableContext, useSortable, verticalListSortingStrategy } from '@dnd-kit/sortable' +import { CSS } from '@dnd-kit/utilities' +import { useVirtualizer } from '@tanstack/react-virtual' +import { type FC, useCallback, useMemo, useRef } from 'react' + +import type { SessionInfo } from '@/hermes' +import { cn } from '@/lib/utils' + +import { SidebarSessionRow } from './session-row' + +interface SessionRowCommonProps { + isPinned: boolean + isSelected: boolean + isWorking: boolean + onDelete: () => void + onPin: () => void + onResume: () => void +} + +interface VirtualSessionListProps { + activeSessionId: null | string + className?: string + onDeleteSession: (sessionId: string) => void + onResumeSession: (sessionId: string) => void + onTogglePin: (sessionId: string) => void + pinned: boolean + sessions: SessionInfo[] + sortable: boolean + workingSessionIdSet: Set +} + +const ROW_ESTIMATE_PX = 28 +const OVERSCAN_ROWS = 12 + +export const VirtualSessionList: FC = ({ + activeSessionId, + className, + onDeleteSession, + onResumeSession, + onTogglePin, + pinned, + sessions, + sortable, + workingSessionIdSet +}) => { + const scrollerRef = useRef(null) + const ids = useMemo(() => sessions.map(s => s.id), [sessions]) + + const virtualizer = useVirtualizer({ + count: sessions.length, + estimateSize: () => ROW_ESTIMATE_PX, + getItemKey: index => sessions[index]?.id ?? index, + getScrollElement: () => scrollerRef.current, + // jsdom-friendly default; the real rect takes over on first observe. + initialRect: { height: 600, width: 240 }, + overscan: OVERSCAN_ROWS + }) + + const virtualItems = virtualizer.getVirtualItems() + const totalSize = virtualizer.getTotalSize() + const paddingTop = virtualItems[0]?.start ?? 0 + const paddingBottom = Math.max(0, totalSize - (virtualItems[virtualItems.length - 1]?.end ?? 0)) + + const rows = virtualItems.map(virtualItem => { + const session = sessions[virtualItem.index] + + if (!session) { + return null + } + + const commonProps: SessionRowCommonProps = { + isPinned: pinned, + isSelected: session.id === activeSessionId, + isWorking: workingSessionIdSet.has(session.id), + onDelete: () => onDeleteSession(session.id), + onPin: () => onTogglePin(session.id), + onResume: () => onResumeSession(session.id) + } + + return sortable ? ( + + ) : ( + + ) + }) + + const list = ( +
+
+ {rows} +
+
+ ) + + return sortable ? ( + + {list} + + ) : ( + list + ) +} + +interface VirtualSortableRowProps { + index: number + measureRef: (node: Element | null) => void + rowProps: SessionRowCommonProps + session: SessionInfo +} + +function VirtualSortableRow({ index, measureRef, rowProps, session }: VirtualSortableRowProps) { + const { attributes, isDragging, listeners, setNodeRef, transform, transition } = useSortable({ id: session.id }) + + // Merge dnd-kit's setNodeRef with the virtualizer's measureElement so + // the row participates in both DnD hit-testing and TanStack height + // measurement. + const refMerged = useCallback( + (node: HTMLDivElement | null) => { + setNodeRef(node) + measureRef(node) + }, + [measureRef, setNodeRef] + ) + + return ( + + ) +} diff --git a/apps/desktop/src/components/assistant-ui/streaming.test.tsx b/apps/desktop/src/components/assistant-ui/streaming.test.tsx index 44e70549083..70f66040e85 100644 --- a/apps/desktop/src/components/assistant-ui/streaming.test.tsx +++ b/apps/desktop/src/components/assistant-ui/streaming.test.tsx @@ -51,6 +51,34 @@ vi.stubGlobal('cancelAnimationFrame', (id: number) => window.clearTimeout(id)) Element.prototype.scrollTo = function scrollTo() {} +Element.prototype.animate = function animate() { + return { + cancel: () => {}, + finished: Promise.resolve() + } as unknown as Animation +} + +// jsdom returns 0 for offset*; the virtualizer reads those to size its +// viewport. Fall through to client* (which tests can override) or a sane +// default so virtualized items render. +function stubOffsetDimension( + prop: 'offsetHeight' | 'offsetWidth', + clientProp: 'clientHeight' | 'clientWidth', + fallback: number +) { + const previous = Object.getOwnPropertyDescriptor(HTMLElement.prototype, prop) + + Object.defineProperty(HTMLElement.prototype, prop, { + configurable: true, + get() { + return previous?.get?.call(this) || (this as HTMLElement)[clientProp] || fallback + } + }) +} + +stubOffsetDimension('offsetWidth', 'clientWidth', 800) +stubOffsetDimension('offsetHeight', 'clientHeight', 600) + async function wait(ms: number) { await act(async () => { await new Promise(resolve => window.setTimeout(resolve, ms)) @@ -85,6 +113,23 @@ function assistantMessage(text: string, running = true): ThreadMessage { } as ThreadMessage } +function assistantErrorMessage(error: string): ThreadMessage { + return { + id: 'assistant-error-1', + role: 'assistant', + content: [], + status: { type: 'incomplete', reason: 'error', error }, + createdAt, + metadata: { + unstable_state: null, + unstable_annotations: [], + unstable_data: [], + steps: [], + custom: {} + } + } as ThreadMessage +} + function assistantReasoningMessage(text: string): ThreadMessage { return { id: 'assistant-reasoning-1', @@ -232,6 +277,20 @@ function TodoHarness({ message }: { message: ThreadMessage }) { ) } +function MessageHarness({ message }: { message: ThreadMessage }) { + const runtime = useExternalStoreRuntime({ + messages: [message], + isRunning: false, + onNew: async () => {} + }) + + return ( + + + + ) +} + function ReasoningHarness() { const runtime = useExternalStoreRuntime({ messages: [assistantReasoningMessage(' The user is asking what this file is.')], @@ -311,6 +370,12 @@ describe('assistant-ui streaming renderer', () => { expect(container.querySelector('[data-slot="aui_composer-clearance"]')).toBeNull() }) + it('renders assistant provider errors inline', () => { + render() + + expect(screen.getByRole('alert').textContent).toContain('OpenRouter rejected the request (403).') + }) + it('does not pull the viewport back down after the user scrolls up during streaming', async () => { const { container } = render() diff --git a/apps/desktop/src/components/assistant-ui/thread-virtualizer.tsx b/apps/desktop/src/components/assistant-ui/thread-virtualizer.tsx new file mode 100644 index 00000000000..bfb7a26aa47 --- /dev/null +++ b/apps/desktop/src/components/assistant-ui/thread-virtualizer.tsx @@ -0,0 +1,306 @@ +import { ThreadPrimitive, useAuiEvent, useAuiState } from '@assistant-ui/react' +import { useVirtualizer, type Virtualizer } from '@tanstack/react-virtual' +import { type ComponentProps, type FC, type ReactNode, useCallback, useEffect, useMemo, useRef } from 'react' + +import { cn } from '@/lib/utils' +import { setThreadScrolledUp } from '@/store/thread-scroll' + +const ESTIMATED_ITEM_HEIGHT = 220 +const OVERSCAN = 4 +const AT_BOTTOM_THRESHOLD = 4 + +type ThreadMessageComponents = ComponentProps['components'] + +type MessageGroup = { id: string; index: number; kind: 'standalone' } | { id: string; indices: number[]; kind: 'turn' } + +interface VirtualizedThreadProps { + clampToComposer: boolean + components: ThreadMessageComponents + emptyPlaceholder?: ReactNode + loadingIndicator?: ReactNode + sessionKey?: string | null +} + +function buildGroups(signature: string): MessageGroup[] { + if (!signature) { + return [] + } + + const messages = signature.split('\n').map(row => { + const [index, id, role] = row.split(':') + + return { id, index: Number(index), role } + }) + + const groups: MessageGroup[] = [] + + for (let i = 0; i < messages.length; i++) { + const message = messages[i] + + if (message.role !== 'user') { + groups.push({ id: message.id, index: message.index, kind: 'standalone' }) + + continue + } + + const indices = [message.index] + + while (i + 1 < messages.length && messages[i + 1].role !== 'user') { + indices.push(messages[++i].index) + } + + groups.push({ id: message.id, indices, kind: 'turn' }) + } + + return groups +} + +export const VirtualizedThread: FC = ({ + clampToComposer, + components, + emptyPlaceholder, + loadingIndicator, + sessionKey +}) => { + const messageSignature = useAuiState(s => + s.thread.messages.map((message, index) => `${index}:${message.id}:${message.role}`).join('\n') + ) + + const groups = useMemo(() => buildGroups(messageSignature), [messageSignature]) + const renderEmpty = groups.length === 0 && Boolean(emptyPlaceholder) + const scrollerRef = useRef(null) + + const virtualizer = useVirtualizer({ + count: groups.length, + estimateSize: () => ESTIMATED_ITEM_HEIGHT, + getItemKey: index => groups[index]?.id ?? index, + getScrollElement: () => scrollerRef.current, + // Seed the rect so the initial range mounts something before + // `observeElementRect` reports the real layout (it overrides this). + initialRect: { height: 600, width: 800 }, + overscan: OVERSCAN + }) + + useThreadScrollAnchor({ + enabled: !renderEmpty, + groupCount: groups.length, + scrollerRef, + sessionKey: sessionKey ?? null, + virtualizer + }) + + const virtualItems = virtualizer.getVirtualItems() + const totalSize = virtualizer.getTotalSize() + const paddingTop = virtualItems[0]?.start ?? 0 + const paddingBottom = Math.max(0, totalSize - (virtualItems.at(-1)?.end ?? 0)) + + return ( +
+
+ {renderEmpty ? ( +
+ {emptyPlaceholder} +
+ ) : ( +
+ {/* Natural-flow virtualization: mounted items render as normal + flex siblings so `position: sticky` on the human bubble + resolves against the scroller without transform interference. + Padding spacers reserve scroll space for unmounted items. */} +
+ {virtualItems.map(virtualItem => { + const group = groups[virtualItem.index] + + if (!group) { + return null + } + + return ( +
+ {group.kind === 'turn' ? ( +
+ {group.indices.map(index => ( + + ))} +
+ ) : ( + + )} +
+ ) + })} +
+ {loadingIndicator} + {clampToComposer && ( + + )} +
+
+ ) +} + +interface ScrollAnchorOptions { + enabled: boolean + groupCount: number + scrollerRef: React.RefObject + sessionKey: string | null + virtualizer: Virtualizer +} + +function useThreadScrollAnchor({ enabled, groupCount, scrollerRef, sessionKey, virtualizer }: ScrollAnchorOptions) { + // `armed` = parked at bottom, content growth should follow. Cleared on + // user-driven upward scroll; re-armed when they reach bottom again. + const armedRef = useRef(true) + const lastTopRef = useRef(0) + const prevSessionKeyRef = useRef(sessionKey) + const prevGroupCountRef = useRef(0) + + const pinToBottom = useCallback(() => { + const el = scrollerRef.current + + if (!el) { + return + } + + el.scrollTop = el.scrollHeight + lastTopRef.current = el.scrollTop + }, [scrollerRef]) + + const jumpToBottom = useCallback(() => { + armedRef.current = true + + if (groupCount > 0) { + virtualizer.scrollToIndex(groupCount - 1, { align: 'end', behavior: 'auto' }) + } + + requestAnimationFrame(() => { + if (armedRef.current) { + pinToBottom() + } + }) + }, [groupCount, pinToBottom, virtualizer]) + + useEffect(() => () => setThreadScrolledUp(false), []) + + // Track at-bottom state, dim composer when scrolled up, disarm on user + // scroll/wheel/touch. + useEffect(() => { + const el = scrollerRef.current + + if (!el) { + return undefined + } + + const disarm = () => { + armedRef.current = false + } + + const onScroll = () => { + const top = el.scrollTop + + if (top + 1 < lastTopRef.current) { + armedRef.current = false + } + + lastTopRef.current = top + + const atBottom = el.scrollHeight - (top + el.clientHeight) <= AT_BOTTOM_THRESHOLD + + if (atBottom) { + armedRef.current = true + } + + setThreadScrolledUp(!atBottom) + } + + const onWheel = (event: WheelEvent) => { + if (event.deltaY < 0) { + disarm() + } + } + + el.addEventListener('scroll', onScroll, { passive: true }) + el.addEventListener('wheel', onWheel, { passive: true }) + el.addEventListener('touchmove', disarm, { passive: true }) + + return () => { + el.removeEventListener('scroll', onScroll) + el.removeEventListener('wheel', onWheel) + el.removeEventListener('touchmove', disarm) + } + }, [scrollerRef]) + + // Follow content growth (streaming, item measurements, loading indicator) + // while armed. + useEffect(() => { + if (!enabled) { + return undefined + } + + const el = scrollerRef.current + + if (!el) { + return undefined + } + + const observer = new ResizeObserver(() => { + if (armedRef.current) { + pinToBottom() + } + }) + + observer.observe(el) + + if (el.firstElementChild) { + observer.observe(el.firstElementChild) + } + + return () => observer.disconnect() + }, [enabled, pinToBottom, scrollerRef]) + + // Jump to bottom on session change OR when an empty thread first gets + // content. Both share the same intent and the same effect. + useEffect(() => { + const sessionChanged = prevSessionKeyRef.current !== sessionKey + const becameNonEmpty = prevGroupCountRef.current === 0 && groupCount > 0 + + prevSessionKeyRef.current = sessionKey + prevGroupCountRef.current = groupCount + + if (enabled && (sessionChanged || becameNonEmpty)) { + jumpToBottom() + } + }, [enabled, groupCount, jumpToBottom, sessionKey]) + + useAuiEvent('thread.runStart', jumpToBottom) +} diff --git a/apps/desktop/src/components/assistant-ui/thread.tsx b/apps/desktop/src/components/assistant-ui/thread.tsx index f93cd542c8c..ecc350d53d9 100644 --- a/apps/desktop/src/components/assistant-ui/thread.tsx +++ b/apps/desktop/src/components/assistant-ui/thread.tsx @@ -1,22 +1,18 @@ import type { Unstable_TriggerAdapter, Unstable_TriggerItem } from '@assistant-ui/core' import { ActionBarPrimitive, - AuiIf, BranchPickerPrimitive, ComposerPrimitive, ErrorPrimitive, MessagePrimitive, - ThreadPrimitive, type ToolCallMessagePartProps, useAui, - useAuiEvent, useAuiState } from '@assistant-ui/react' import { useStore } from '@nanostores/react' import { IconPlayerStopFilled } from '@tabler/icons-react' import { type ClipboardEvent, - type ComponentProps, type FC, type FocusEvent, type FormEvent, @@ -29,7 +25,6 @@ import { useRef, useState } from 'react' -import { StickToBottom, useStickToBottomContext } from 'use-stick-to-bottom' import { COMPOSER_DROP_ACTIVE_CLASS, COMPOSER_DROP_FADE_CLASS } from '@/app/chat/composer/drop-affordance' import { @@ -56,6 +51,7 @@ import { ClarifyTool } from '@/components/assistant-ui/clarify-tool' import { DirectiveContent, DirectiveText } from '@/components/assistant-ui/directive-text' import { hermesDirectiveFormatter } from '@/components/assistant-ui/directive-text' import { MarkdownText } from '@/components/assistant-ui/markdown-text' +import { VirtualizedThread } from '@/components/assistant-ui/thread-virtualizer' import { HoistedTodoPanel, todosFromMessageContent } from '@/components/assistant-ui/todo-tool' import { ToolFallback, ToolGroupSlot } from '@/components/assistant-ui/tool-fallback' import { TooltipIconButton } from '@/components/assistant-ui/tooltip-icon-button' @@ -85,16 +81,10 @@ import { useEnterAnimation } from '@/lib/use-enter-animation' import { cn } from '@/lib/utils' import { playSpeechText, stopVoicePlayback } from '@/lib/voice-playback' import { notifyError } from '@/store/notifications' -import { setThreadScrolledUp } from '@/store/thread-scroll' import { $voicePlayback } from '@/store/voice-playback' type ThreadLoadingState = 'response' | 'session' -interface StickyStateFlags { - escapedFromLock: boolean - isAtBottom: boolean -} - interface MessageActionProps { messageId: string messageText: string @@ -129,17 +119,6 @@ const INTERRUPTED_ONLY_RE = /^_?\[interrupted\]_?$/i const isInterruptedOnlyMessage = (text: string) => INTERRUPTED_ONLY_RE.test(text.trim()) -function resetStickyState(state: StickyStateFlags) { - state.escapedFromLock = false - state.isAtBottom = true -} - -function pinElementToBottom(el: HTMLElement) { - el.scrollTop = el.scrollHeight - - return el.scrollTop -} - export const Thread: FC<{ clampToComposer?: boolean cwd?: string | null @@ -161,8 +140,6 @@ export const Thread: FC<{ sessionId = null, sessionKey }) => { - const introHero = useAuiState(s => Boolean(intro) && s.thread.isEmpty) - const messageComponents = useMemo( () => ({ AssistantMessage: () => , @@ -173,279 +150,31 @@ export const Thread: FC<{ [cwd, gateway, onBranchInNewChat, onCancel, sessionId] ) + const emptyPlaceholder = intro ? ( +
+ +
+ ) : undefined + return ( - - - - - - Boolean(intro) && s.thread.isEmpty}> - {intro ? ( -
- -
- ) : null} - - - {loading === 'response' && } - {clampToComposer && ( -