Skip to content

Virtualize thread messages with react-virtual #389

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
urjitc merged 2 commits into from
Apr 17, 2026
Merged

Virtualize thread messages with react-virtual #389

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d060951
Virtualize thread messages with @tanstack/react-virtual to fix 1.2s k…
urjitc Apr 17, 2026
c309242
Bump estimateSize to 350 and use virtualRow.key for React reconciliation
urjitc Apr 17, 2026
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
1 change: 1 addition & 0 deletions package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@
"@tanstack/react-pacer": "^0.21.1",
"@tanstack/react-query": "^5.96.1",
"@tanstack/react-query-devtools": "^5.96.1",
"@tanstack/react-virtual": "^3.13.23",
"@tiptap/core": "3.22.3",
"@tiptap/extension-code-block": "3.22.3",
"@tiptap/extension-highlight": "3.22.3",
Expand Down
72 changes: 61 additions & 11 deletions src/components/assistant-ui/thread.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,9 @@ import {
useMessagePartText,
useAuiState,
} from "@assistant-ui/react";
import { useVirtualizer } from "@tanstack/react-virtual";

import type { FC } from "react";
import type { FC, RefObject } from "react";
import { createContext, useContext } from "react";
import { useEffect, useRef, useState, useMemo, useCallback } from "react";

Expand Down Expand Up @@ -137,13 +138,7 @@ export const Thread: FC<ThreadProps> = ({ items = [] }) => {
<ThreadWelcome items={items} />
</AuiIf>

<ThreadPrimitive.Messages
components={{
UserMessage,
EditComposer,
AssistantMessage,
}}
/>
<VirtualizedMessages scrollRef={viewportRef} />

@coderabbitai coderabbitai Bot Apr 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Accessibility / find-in-page regression from virtualization.

Only rows within the overscan window are in the DOM, so browser Find (Ctrl/Cmd+F), "scroll to selection", copy-across-messages, and screen-reader sequential navigation no longer cover the full thread. For a chat transcript this is a meaningful UX/a11y regression worth calling out in the PR. Consider mitigations such as: increasing overscan significantly, rendering an off-screen plain-text summary for AT, or re-mounting all rows when the user triggers in-page search (e.g., via a "show full thread" affordance). At minimum, document the tradeoff.

Also applies to: 155-196

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/components/assistant-ui/thread.tsx` at line 141, The VirtualizedMessages
component currently only renders rows inside the overscan window (used in Thread
via <VirtualizedMessages scrollRef={viewportRef} />), which breaks browser Find,
selection-copy across messages, and screen-reader navigation; fix by one of: (a)
increase the virtualization overscan prop on VirtualizedMessages to a much
larger value (set overscan to a high number or time-based buffer) so more items
are in DOM, (b) add an off-screen plain-text transcript renderer inside the
Thread component (hidden visually but available to AT) that outputs the full
thread as selectable text for find/copy/AT, or (c) add a “show full thread”
affordance that remounts VirtualizedMessages to render all rows (disable
virtualization) when the user triggers in-page search or clicks the affordance;
implement the chosen mitigation in Thread/VirtualizedMessages and add a short
code comment documenting the tradeoff if you opt to keep virtualization.

</ThreadPrimitive.Viewport>

<div className="aui-thread-composer-wrapper mx-auto flex w-full max-w-[var(--thread-max-width)] flex-shrink-0 flex-col gap-4 overflow-visible rounded-t-3xl bg-sidebar px-4 pb-3 md:pb-4">
Expand All @@ -153,6 +148,53 @@ export const Thread: FC<ThreadProps> = ({ items = [] }) => {
);
};

interface VirtualizedMessagesProps {
scrollRef: RefObject<HTMLDivElement | null>;
}

const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({ scrollRef }) => {
const messageCount = useAuiState((s) => s.thread.messages.length);

const virtualizer = useVirtualizer({
count: messageCount,
getScrollElement: () => scrollRef.current,
estimateSize: () => 350,
overscan: 5,

@greptile-apps greptile-apps Bot Apr 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 estimateSize of 150 is likely too low for assistant messages

Chat messages span a wide range — user queries can be a single line (~50 px) while assistant responses with markdown, code blocks, and reasoning sections can easily exceed 1 000 px. With a 150 px estimate, the virtual container's initial height for a 100-message thread is only 15 000 px; as measureElement corrects taller items, the container jumps in height, causing visible layout shifts when scrolling through older parts of a thread for the first time. A higher default (e.g. 300–400 px) better approximates the average assistant message and reduces those jumps.

Suggested change
overscan: 5,
estimateSize: () => 350,

Fix in Cursor

});

if (messageCount === 0) return null;

return (
<div
style={{
height: `${virtualizer.getTotalSize()}px`,
width: "100%",
position: "relative",
}}
>
{virtualizer.getVirtualItems().map((virtualRow) => (
Comment thread
capy-ai[bot] marked this conversation as resolved.
<div
key={virtualRow.key}
data-index={virtualRow.index}
ref={virtualizer.measureElement}
style={{
position: "absolute",
top: 0,
left: 0,
width: "100%",
transform: `translateY(${virtualRow.start}px)`,
}}
>
<ThreadPrimitive.MessageByIndex
index={virtualRow.index}
components={MESSAGE_COMPONENTS}
/>
</div>
))}
</div>
);
};
Comment on lines +155 to +196

@coderabbitai coderabbitai Bot Apr 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🌐 Web query:

@tanstack/react-virtual getItemKey usage with measureElement for dynamic lists

💡 Result:

In @tanstack/react-virtual, getItemKey is an optional callback passed to useVirtualizer that returns a unique identifier (string, number, or bigint) for each item based on its index. By default, it returns the index itself. Use it for dynamic lists with stable unique IDs (e.g., database IDs) to improve React reconciliation, caching, and prevent unnecessary re-renders/measures when items are added/removed/reordered. It works seamlessly with measureElement for dynamic sizing: 1. Provide getItemKey in virtualizer options (memoize with useCallback to avoid infinite re-renders): const getItemKey = useCallback((index: number) => items[index]?.id ?? index, [items]); const virtualizer = useVirtualizer({ count: items.length, getScrollElement: => parentRef.current, estimateSize: => 50, measureElement: (element) => element ? element.getBoundingClientRect.height : 0, getItemKey, }); 2. In the render loop, use virtualItem.key as the React key prop and virtualizer.measureElement as ref: {virtualizer.getVirtualItems.map(virtualItem => ( <div key={virtualItem.key} // Uses getItemKey-derived key data-index={virtualItem.index} ref={virtualizer.measureElement} style={{ position: 'absolute', top: 0, left: 0, width: '100%', transform: translateY(${virtualItem.start}px), height: virtualItem.size, // Updates after measurement }} > {items[virtualItem.index]} ))} Key benefits for dynamic lists: - Caches measurements by unique key (not index), surviving list mutations. - Stabilizes React keys for better diffing. - Official docs recommend overriding with unique IDs "when possible". Pitfalls (from GitHub issues): - Must memoize getItemKey (e.g., useCallback)—unmemoized causes cache invalidation/re-renders. - Ensure getItemKey dependencies (e.g., items) are stable. - Use virtualItem.key (not items[virtualItem.index].id) for React key to match virtualizer's internal mapping. - Callback refs like ref={(el) => virtualizer.measureElement(el)} may interfere with caching; prefer direct ref={virtualizer.measureElement}. Official dynamic example uses default virtualRow.key (index-based) successfully, but custom keys enhance performance for real-world dynamic data. Examples: Official dynamic demo (tanstack.com/virtual/latest/docs/framework/react/examples/dynamic), LogRocket chat tutorial (uses item.id as React key with measureElement).

Citations:

Provide a stable getItemKey so measurement cache stays correct across message changes.

Rows are keyed by virtualRow.index and measureElement caches measured heights by that index. When messages are added/removed, or when branches are switched, the index→message mapping shifts and cached heights end up attached to the wrong message, causing layout jumps and incorrect getTotalSize() until every visible row re-measures. Pass a stable ID via getItemKey and use virtualRow.key as the React key prop.

♻️ Suggested fix 
 const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({ scrollRef }) => {
   const messageCount = useAuiState((s) => s.thread.messages.length);
+  const getMessageId = useCallback(
+    (index: number) => useAuiState.getState().thread.messages[index]?.id ?? index,
+    []
+  );

   const virtualizer = useVirtualizer({
     count: messageCount,
     getScrollElement: () => scrollRef.current,
     estimateSize: () => 150,
     overscan: 5,
+    getItemKey: (index) => getMessageId(index),
   });
   ...
       {virtualizer.getVirtualItems().map((virtualRow) => (
         <div
-          key={virtualRow.index}
+          key={virtualRow.key}
           data-index={virtualRow.index}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({ scrollRef }) => {
const messageCount = useAuiState((s) => s.thread.messages.length);
const virtualizer = useVirtualizer({
count: messageCount,
getScrollElement: () => scrollRef.current,
estimateSize: () => 150,
overscan: 5,
});
if (messageCount === 0) return null;
return (
<div
style={{
height: `${virtualizer.getTotalSize()}px`,
width: "100%",
position: "relative",
}}
>
{virtualizer.getVirtualItems().map((virtualRow) => (
<div
key={virtualRow.index}
data-index={virtualRow.index}
ref={virtualizer.measureElement}
style={{
position: "absolute",
top: 0,
left: 0,
width: "100%",
transform: `translateY(${virtualRow.start}px)`,
}}
>
<ThreadPrimitive.MessageByIndex
index={virtualRow.index}
components={MESSAGE_COMPONENTS}
/>
</div>
))}
</div>
);
};
const VirtualizedMessages: FC<VirtualizedMessagesProps> = ({ scrollRef }) => {
const messageCount = useAuiState((s) => s.thread.messages.length);
const getMessageId = useCallback(
(index: number) => useAuiState.getState().thread.messages[index]?.id ?? index,
[]
);
const virtualizer = useVirtualizer({
count: messageCount,
getScrollElement: () => scrollRef.current,
estimateSize: () => 150,
overscan: 5,
getItemKey: (index) => getMessageId(index),
});
if (messageCount === 0) return null;
return (
<div
style={{
height: `${virtualizer.getTotalSize()}px`,
width: "100%",
position: "relative",
}}
>
{virtualizer.getVirtualItems().map((virtualRow) => (
<div
key={virtualRow.key}
data-index={virtualRow.index}
ref={virtualizer.measureElement}
style={{
position: "absolute",
top: 0,
left: 0,
width: "100%",
transform: `translateY(${virtualRow.start}px)`,
}}
>
<ThreadPrimitive.MessageByIndex
index={virtualRow.index}
components={MESSAGE_COMPONENTS}
/>
</div>
))}
</div>
);
};
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/components/assistant-ui/thread.tsx` around lines 155 - 196, The
virtualizer is currently keyed by virtualRow.index which shifts when messages
change causing measurement cache corruption; update the useVirtualizer call in
VirtualizedMessages to provide a stable getItemKey that returns a unique message
id for a given index (lookup from the same thread/messages state used to compute
messageCount), and then use virtualRow.key (not virtualRow.index) as the React
key for the rendered row that calls virtualizer.measureElement and renders
ThreadPrimitive.MessageByIndex with MESSAGE_COMPONENTS so measured heights
remain attached to the correct message across inserts/removals/branch switches.

Comment on lines +155 to +196

@greptile-apps greptile-apps Bot Apr 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Verify streaming scroll behavior is preserved

ThreadPrimitive.Messages may have contained its own scroll-to-bottom logic that fired as the AI streamed tokens, independent of ThreadPrimitive.Viewport's autoScroll flag. The new VirtualizedMessages uses the virtualizer's raw scroll APIs and has no useEffect that calls virtualizer.scrollToIndex(messageCount - 1) when messageCount increases. If the old component did perform any self-managed "follow the stream" scrolling, that behavior is now gone. Worth confirming manually that, during an active generation, a user who was at the bottom of the thread continues to see new tokens appear without having to scroll.

Fix in Cursor


const ThreadLoadingSkeleton: FC = () => {
return (
<div
Expand Down Expand Up @@ -396,10 +438,12 @@ const ComposerHoverWrapper: FC<ComposerHoverWrapperProps> = ({ items }) => {
null,
);
const isThreadEmpty = useAuiState(({ thread }) => thread?.isEmpty ?? true);
const composerText = useAuiState(
(s) => (s as { composer?: { text?: string } })?.composer?.text ?? "",
const hasComposerText = useAuiState(
(s) =>
Boolean(
((s as { composer?: { text?: string } })?.composer?.text ?? "").trim(),
),
);
const hasComposerText = Boolean(composerText?.trim());

const handleDirectFill = useCallback(
(fill: string) => {
Expand Down Expand Up @@ -1429,6 +1473,12 @@ const EditComposer: FC = () => {
);
};

const MESSAGE_COMPONENTS = {
UserMessage,
EditComposer,
AssistantMessage,
};

const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({
className,
...rest
Expand Down
Loading