From 5aaf4cce36c9c215819345302e2d9e3194c08ca7 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Fri, 17 Apr 2026 21:56:02 +0000 Subject: [PATCH 01/12] Replace assistant-ui Viewport with custom scroll controller to resolve virtualizer conflicts --- src/components/assistant-ui/thread.tsx | 318 ++++++++++++++++++------- 1 file changed, 237 insertions(+), 81 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index e2efadc8..bfeecd6b 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -35,13 +35,14 @@ import { MessagePrimitive, ThreadPrimitive, useAui, + useAuiEvent, useMessage, useMessagePartText, useAuiState, } from "@assistant-ui/react"; import { useVirtualizer } from "@tanstack/react-virtual"; -import type { FC, RefObject } from "react"; +import type { FC, MutableRefObject, RefObject } from "react"; import { createContext, useContext } from "react"; import { useEffect, useRef, useState, useMemo, useCallback } from "react"; @@ -117,6 +118,7 @@ interface ThreadProps { export const Thread: FC = ({ items = [] }) => { const viewportRef = useRef(null); + const virtualizerHandleRef = useRef(null); return ( = ({ items = [] }) => { ["--thread-max-width" as string]: "50rem", }} > - thread.isLoading}> @@ -138,8 +147,15 @@ export const Thread: FC = ({ items = [] }) => { - - + + +
@@ -148,20 +164,154 @@ export const Thread: FC = ({ items = [] }) => { ); }; +/** + * Imperative handle exposed by VirtualizedMessages so ThreadScrollController + * can translate message indices into scrollTop offsets via @tanstack/react-virtual. + */ +interface VirtualizerHandle { + /** Returns the virtualRow.start for the given message index, or null if not currently mounted/measured. */ + getRowOffset: (index: number) => number | null; + /** Returns current message count. */ + getCount: () => number; + /** Returns the estimateSize() value used by the virtualizer (fallback for unmounted rows). */ + getEstimatedSize: () => number; +} + +interface ThreadScrollControllerProps { + scrollRef: RefObject; + virtualizerHandleRef: MutableRefObject; +} + +/** + * Owns scroll behavior for the virtualized thread. Replaces assistant-ui's + * ThreadPrimitive.Viewport + turnAnchor + ViewportSlack entirely. + * + * Behavior: + * - thread.initialize (history first load) → bottom, instant. + * - threadListItem.switchedTo (thread switch) → bottom, instant. + * - thread.runStart (new user message sent) → anchor that user message at the top + * of the viewport, smooth. The "new user message" is messages.length - 1 at runStart. + * - Streaming: no auto-scroll. User owns scroll. + * + * We never install ResizeObserver/MutationObserver on the scroll container — + * that loop is exactly what was fighting the virtualizer. + */ +const ThreadScrollController: FC = ({ + scrollRef, + virtualizerHandleRef, +}) => { + const scrollToBottom = useCallback( + (behavior: ScrollBehavior) => { + requestAnimationFrame(() => { + const el = scrollRef.current; + if (!el) return; + el.scrollTo({ top: el.scrollHeight, behavior }); + }); + }, + [scrollRef], + ); + + const scrollIndexToTop = useCallback( + (targetIndex: number, behavior: ScrollBehavior) => { + // React-virtual may not have measured the new row yet. Retry across a few + // frames using the measured offset when available, falling back to + // estimateSize * index so we at least land close. + let attempts = 0; + const attempt = () => { + attempts += 1; + const el = scrollRef.current; + const handle = virtualizerHandleRef.current; + if (!el || !handle) return; + if (targetIndex < 0 || targetIndex >= handle.getCount()) return; + + const measured = handle.getRowOffset(targetIndex); + const top = measured ?? targetIndex * handle.getEstimatedSize(); + el.scrollTo({ top, behavior }); + + if (measured === null && attempts < 3) { + requestAnimationFrame(attempt); + } + }; + requestAnimationFrame(attempt); + }, + [scrollRef, virtualizerHandleRef], + ); + + useAuiEvent("thread.initialize", () => { + scrollToBottom("instant"); + }); + + useAuiEvent("threadListItem.switchedTo", () => { + scrollToBottom("instant"); + }); + + useAuiEvent("thread.runStart", () => { + const handle = virtualizerHandleRef.current; + if (!handle) return; + // At runStart, the just-sent user message is at messages.length - 1. + // The assistant placeholder is appended shortly after, so when the + // assistant row arrives our user message will naturally be one up. + const targetIndex = handle.getCount() - 1; + scrollIndexToTop(targetIndex, "smooth"); + }); + + return null; +}; + +/** + * Replacement for the hover tracking that `MessagePrimitive.Root` normally wires + * up (via ThreadPrimitiveViewportSlack). Without this, ActionBarPrimitive + * autohide="not-last" and MessagePrimitive.If lastOrHover would be permanently + * hidden on non-last messages. We forward hover state into the aui message + * client directly, which is what MessageRoot does under the hood. + */ +const useMessageHoverHandlers = () => { + const aui = useAui(); + const onMouseEnter = useCallback(() => { + aui?.message()?.setIsHovering?.(true); + }, [aui]); + const onMouseLeave = useCallback(() => { + aui?.message()?.setIsHovering?.(false); + }, [aui]); + return { onMouseEnter, onMouseLeave }; +}; + interface VirtualizedMessagesProps { scrollRef: RefObject; + virtualizerHandleRef: MutableRefObject; } -const VirtualizedMessages: FC = ({ scrollRef }) => { +const ROW_ESTIMATE_SIZE = 350; + +const VirtualizedMessages: FC = ({ + scrollRef, + virtualizerHandleRef, +}) => { const messageCount = useAuiState((s) => s.thread.messages.length); const virtualizer = useVirtualizer({ count: messageCount, getScrollElement: () => scrollRef.current, - estimateSize: () => 350, + estimateSize: () => ROW_ESTIMATE_SIZE, overscan: 5, }); + // Publish imperative handle for ThreadScrollController. + useEffect(() => { + virtualizerHandleRef.current = { + getRowOffset: (index) => { + const items = virtualizer.getVirtualItems(); + const found = items.find((v) => v.index === index); + return found ? found.start : null; + }, + getCount: () => messageCount, + getEstimatedSize: () => ROW_ESTIMATE_SIZE, + }; + return () => { + virtualizerHandleRef.current = null; + }; + }, [virtualizer, messageCount, virtualizerHandleRef]); + if (messageCount === 0) return null; return ( @@ -1185,37 +1335,40 @@ const MessageError: FC = () => { }; const AssistantMessage: FC = () => { + const message = useMessage(); + const hover = useMessageHoverHandlers(); return ( - -
-
- - - -
+
+
+ + + +
-
- - -
+
+ +
- +
); }; @@ -1311,55 +1464,58 @@ const UserMessage: FC = () => { [expanded, showExpand], ); - return ( - -
- {/* Attachments display */} - - -
- - -
- - {showExpand && ( -
- -
- )} -
- -
+ const hover = useMessageHoverHandlers(); -
-
- + return ( +
+ {/* Attachments display */} + + +
+ + +
+ + {showExpand && ( +
+ +
+ )}
-
+ +
- +
+
+ +
- + + +
); }; From 4083551eee29590994f5e9f0523cb2b34f030d40 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:04:55 +0000 Subject: [PATCH 02/12] fix(thread): use virtualizer.scrollToIndex + live message count Addresses P2 review comments from Greptile and Cubic: 1. Replace hand-rolled offset math (getRowOffset + index*350 fallback + RAF retry loop) with TanStack Virtual's own scrollToIndex. That already handles unmeasured rows via its internal size ledger and retries up to 10 frames, so if the user scrolled up before sending, the target now uses actual measured heights instead of the estimateSize fallback. 2. Read the live message count from aui.thread().getState().messages at thread.runStart time instead of through the VirtualizedMessages effect closure. That closure could return the pre-send count if react effect ordering hadn't committed yet, causing us to scroll to the previous message. --- src/components/assistant-ui/thread.tsx | 83 +++++++++++--------------- 1 file changed, 36 insertions(+), 47 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index bfeecd6b..6e766aeb 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -166,15 +166,20 @@ export const Thread: FC = ({ items = [] }) => { /** * Imperative handle exposed by VirtualizedMessages so ThreadScrollController - * can translate message indices into scrollTop offsets via @tanstack/react-virtual. + * can position the viewport on a specific message index. We delegate to + * TanStack Virtual's own `scrollToIndex` so unmeasured rows use the internal + * size ledger (instead of reimplementing offset math ourselves). */ interface VirtualizerHandle { - /** Returns the virtualRow.start for the given message index, or null if not currently mounted/measured. */ - getRowOffset: (index: number) => number | null; - /** Returns current message count. */ - getCount: () => number; - /** Returns the estimateSize() value used by the virtualizer (fallback for unmounted rows). */ - getEstimatedSize: () => number; + /** + * Scroll so the row at `index` aligns according to `align`. + * Uses virtualizer.scrollToIndex internally, which retries across frames + * for rows whose size hasn't been measured yet. + */ + scrollToIndex: ( + index: number, + options?: { align?: "start" | "center" | "end" | "auto"; behavior?: ScrollBehavior }, + ) => void; } interface ThreadScrollControllerProps { @@ -200,6 +205,8 @@ const ThreadScrollController: FC = ({ scrollRef, virtualizerHandleRef, }) => { + const aui = useAui(); + const scrollToBottom = useCallback( (behavior: ScrollBehavior) => { requestAnimationFrame(() => { @@ -211,32 +218,6 @@ const ThreadScrollController: FC = ({ [scrollRef], ); - const scrollIndexToTop = useCallback( - (targetIndex: number, behavior: ScrollBehavior) => { - // React-virtual may not have measured the new row yet. Retry across a few - // frames using the measured offset when available, falling back to - // estimateSize * index so we at least land close. - let attempts = 0; - const attempt = () => { - attempts += 1; - const el = scrollRef.current; - const handle = virtualizerHandleRef.current; - if (!el || !handle) return; - if (targetIndex < 0 || targetIndex >= handle.getCount()) return; - - const measured = handle.getRowOffset(targetIndex); - const top = measured ?? targetIndex * handle.getEstimatedSize(); - el.scrollTo({ top, behavior }); - - if (measured === null && attempts < 3) { - requestAnimationFrame(attempt); - } - }; - requestAnimationFrame(attempt); - }, - [scrollRef, virtualizerHandleRef], - ); - useAuiEvent("thread.initialize", () => { scrollToBottom("instant"); }); @@ -248,11 +229,24 @@ const ThreadScrollController: FC = ({ useAuiEvent("thread.runStart", () => { const handle = virtualizerHandleRef.current; if (!handle) return; - // At runStart, the just-sent user message is at messages.length - 1. - // The assistant placeholder is appended shortly after, so when the - // assistant row arrives our user message will naturally be one up. - const targetIndex = handle.getCount() - 1; - scrollIndexToTop(targetIndex, "smooth"); + // Read the live message count from the aui thread state at event time. + // Going through the handle's closure over `messageCount` is unsafe because + // React sibling effects may not have committed the new count yet when the + // event fires — we'd then scroll to the previous message. + const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; + if (liveCount <= 0) return; + // The just-sent user message is at messages.length - 1 at runStart time. + // When the assistant row arrives the user message naturally ends up one up. + const targetIndex = liveCount - 1; + // scrollToIndex uses the virtualizer's internal size ledger (including + // previously measured heights) and retries internally for unmeasured rows. + // Wrap in rAF so react-virtual sees the new row count before we scroll. + requestAnimationFrame(() => { + virtualizerHandleRef.current?.scrollToIndex(targetIndex, { + align: "start", + behavior: "smooth", + }); + }); }); return null; @@ -296,21 +290,16 @@ const VirtualizedMessages: FC = ({ overscan: 5, }); - // Publish imperative handle for ThreadScrollController. + // Publish imperative handle for ThreadScrollController. We expose the + // virtualizer's own scrollToIndex so callers don't have to recompute offsets. useEffect(() => { virtualizerHandleRef.current = { - getRowOffset: (index) => { - const items = virtualizer.getVirtualItems(); - const found = items.find((v) => v.index === index); - return found ? found.start : null; - }, - getCount: () => messageCount, - getEstimatedSize: () => ROW_ESTIMATE_SIZE, + scrollToIndex: (index, options) => virtualizer.scrollToIndex(index, options), }; return () => { virtualizerHandleRef.current = null; }; - }, [virtualizer, messageCount, virtualizerHandleRef]); + }, [virtualizer, virtualizerHandleRef]); if (messageCount === 0) return null; From 9dd0382e48515da89000062cadef4f8886d3885f Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:15:37 +0000 Subject: [PATCH 03/12] perf(thread): align virtualized thread with TanStack Virtual docs Full refactor of the virtualized thread following TanStack Virtual official patterns (docs/framework/react/examples/dynamic): Scroll container (the outer viewport): - Add `contain: strict` so the browser treats the thread as an independent layout/paint root, cutting reflow cost during streaming. - Add `overflow-anchor: none` to disable native browser scroll anchoring, which otherwise fights the virtualizer when row heights change mid-stream. - Add `scrollPaddingTop` matching the virtualizer's scrollPaddingStart so keyboard scrolling / scroll-into-view land with breathing room. Virtualizer options: - `getItemKey: (index) => messageIds[index]` \u2014 stable keys prevent measurementsCache invalidation when branches flip or messages reorder. Previously keys were implicit indices and a branch switch would throw away all measurements. - `scrollPaddingStart: SCROLL_PADDING_START` so scrollToIndex(..., { align: 'start' }) leaves room above the user message anchor. - `useScrollendEvent: true` \u2014 use the native scrollend event on modern browsers instead of the virtualizer's polling fallback. - `enabled: mounted` \u2014 don't compute a range on SSR / first render against a null scroll element. - `rangeExtractor` \u2014 union the default range with the streaming last assistant message index while thread.isRunning, so scrolling away during a stream doesn't unmount the message (which was killing live useMessagePartText updates and hover state). Render shape (docs pattern): - Single absolutely-positioned wrapper translated once by `items[0]?.start ?? 0`. - Rows inside render in natural document order (no per-row transform). That's noticeably cheaper than the previous per-row translateY and plays nicely with contain:strict on the parent. Scroll controller cleanup: - Dropped the VirtualizerHandle abstraction. The virtualizer instance is lifted into Thread via a ref and ThreadScrollController calls virtualizer.scrollToIndex directly. - All three scroll entry points (thread.initialize, switchedTo, runStart) share a single scrollToLastIndex(align, behavior) helper that reads the live message count from aui.thread().getState() at event time (safe against sibling-effect ordering). --- src/components/assistant-ui/thread.tsx | 244 +++++++++++++++---------- 1 file changed, 147 insertions(+), 97 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 6e766aeb..7d60f549 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -40,7 +40,11 @@ import { useMessagePartText, useAuiState, } from "@assistant-ui/react"; -import { useVirtualizer } from "@tanstack/react-virtual"; +import { + defaultRangeExtractor, + useVirtualizer, + type Virtualizer, +} from "@tanstack/react-virtual"; import type { FC, MutableRefObject, RefObject } from "react"; import { createContext, useContext } from "react"; @@ -116,9 +120,35 @@ interface ThreadProps { items?: Item[]; } +// Shared type alias: the @tanstack/react-virtual instance used throughout this +// file. We lift it to Thread so the scroll controller can call scrollToIndex +// directly without going through a handle wrapper. +type ThreadVirtualizer = Virtualizer; + +// Estimated row height used before the virtualizer measures a row's actual +// size. Matches the previous value — tuned for a typical assistant turn with +// some markdown and a tool call. +const ROW_ESTIMATE_SIZE = 350; +// Breathing room above a user message when we scrollToIndex(..., { align: 'start' }). +// Without this the message is flush against the viewport's top edge, which +// looks wrong and hides MessageContextBadges on touch. +const SCROLL_PADDING_START = 16; +// Rows kept mounted outside the visible viewport (above and below) so smooth +// scrolling doesn't flash blank regions. 5 matches TanStack's own docs example. +const OVERSCAN = 5; + export const Thread: FC = ({ items = [] }) => { const viewportRef = useRef(null); - const virtualizerHandleRef = useRef(null); + // Lifted so ThreadScrollController can imperatively scroll to an index + // without a separate "handle" abstraction. + const virtualizerRef = useRef(null); + + // Gate useVirtualizer behind mount so SSR/first-paint doesn't compute a + // range against a null scroll element (TanStack's `enabled` option). + const [mounted, setMounted] = useState(false); + useEffect(() => { + setMounted(true); + }, []); return ( = ({ items = [] }) => { assistant-ui's Viewport installs a ResizeObserver + subtree MutationObserver on the scroll container and MessagePrimitive.Root injects inline min-height on the last assistant message via ThreadViewportSlack. Both fight our - @tanstack/react-virtual virtualizer (observers fire constantly, Slack's - injected min-height is read by measureElement as the row size). We own - scroll behavior via ThreadScrollController below. + @tanstack/react-virtual virtualizer. We own scroll behavior here. + + - `contain: strict` lets the browser treat this subtree as an independent + layout/paint root, cutting reflow cost during streaming. + - `overflow-anchor: none` disables the browser's native scroll-anchoring, + which otherwise fights the virtualizer when row sizes change. */}
thread.isLoading}> @@ -149,12 +187,10 @@ export const Thread: FC = ({ items = [] }) => { - +
@@ -164,27 +200,8 @@ export const Thread: FC = ({ items = [] }) => { ); }; -/** - * Imperative handle exposed by VirtualizedMessages so ThreadScrollController - * can position the viewport on a specific message index. We delegate to - * TanStack Virtual's own `scrollToIndex` so unmeasured rows use the internal - * size ledger (instead of reimplementing offset math ourselves). - */ -interface VirtualizerHandle { - /** - * Scroll so the row at `index` aligns according to `align`. - * Uses virtualizer.scrollToIndex internally, which retries across frames - * for rows whose size hasn't been measured yet. - */ - scrollToIndex: ( - index: number, - options?: { align?: "start" | "center" | "end" | "auto"; behavior?: ScrollBehavior }, - ) => void; -} - interface ThreadScrollControllerProps { - scrollRef: RefObject; - virtualizerHandleRef: MutableRefObject; + virtualizerRef: MutableRefObject; } /** @@ -194,59 +211,50 @@ interface ThreadScrollControllerProps { * Behavior: * - thread.initialize (history first load) → bottom, instant. * - threadListItem.switchedTo (thread switch) → bottom, instant. - * - thread.runStart (new user message sent) → anchor that user message at the top - * of the viewport, smooth. The "new user message" is messages.length - 1 at runStart. + * - thread.runStart (new user message sent) → anchor that user message at the + * top of the viewport with scrollPaddingStart breathing room, smooth. * - Streaming: no auto-scroll. User owns scroll. * - * We never install ResizeObserver/MutationObserver on the scroll container — - * that loop is exactly what was fighting the virtualizer. + * Everything goes through virtualizer.scrollToIndex so unmeasured rows use + * the internal size ledger and built-in retry logic. */ const ThreadScrollController: FC = ({ - scrollRef, - virtualizerHandleRef, + virtualizerRef, }) => { const aui = useAui(); - const scrollToBottom = useCallback( - (behavior: ScrollBehavior) => { + const scrollToLastIndex = useCallback( + (align: "start" | "end", behavior: ScrollBehavior) => { + const count = aui?.thread()?.getState()?.messages.length ?? 0; + if (count <= 0) return; + // Read the live count from aui state instead of any closure so we're + // safe against sibling effect ordering (VirtualizedMessages may not + // have committed the new count when this event fires). + const targetIndex = count - 1; + // rAF lets react-virtual see the new row count in its next commit + // before we scroll; scrollToIndex itself then retries internally for + // rows whose size hasn't been measured yet. requestAnimationFrame(() => { - const el = scrollRef.current; - if (!el) return; - el.scrollTo({ top: el.scrollHeight, behavior }); + virtualizerRef.current?.scrollToIndex(targetIndex, { align, behavior }); }); }, - [scrollRef], + [aui, virtualizerRef], ); useAuiEvent("thread.initialize", () => { - scrollToBottom("instant"); + scrollToLastIndex("end", "instant"); }); useAuiEvent("threadListItem.switchedTo", () => { - scrollToBottom("instant"); + scrollToLastIndex("end", "instant"); }); useAuiEvent("thread.runStart", () => { - const handle = virtualizerHandleRef.current; - if (!handle) return; - // Read the live message count from the aui thread state at event time. - // Going through the handle's closure over `messageCount` is unsafe because - // React sibling effects may not have committed the new count yet when the - // event fires — we'd then scroll to the previous message. - const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; - if (liveCount <= 0) return; - // The just-sent user message is at messages.length - 1 at runStart time. - // When the assistant row arrives the user message naturally ends up one up. - const targetIndex = liveCount - 1; - // scrollToIndex uses the virtualizer's internal size ledger (including - // previously measured heights) and retries internally for unmeasured rows. - // Wrap in rAF so react-virtual sees the new row count before we scroll. - requestAnimationFrame(() => { - virtualizerHandleRef.current?.scrollToIndex(targetIndex, { - align: "start", - behavior: "smooth", - }); - }); + // At runStart the just-sent user message is at messages.length - 1. + // Anchoring it at the top mirrors the top-turnAnchor feel we previously + // got from assistant-ui, without the min-height injection that was + // fighting our virtualizer. + scrollToLastIndex("start", "smooth"); }); return null; @@ -272,37 +280,76 @@ const useMessageHoverHandlers = () => { interface VirtualizedMessagesProps { scrollRef: RefObject; - virtualizerHandleRef: MutableRefObject; + virtualizerRef: MutableRefObject; + enabled: boolean; } -const ROW_ESTIMATE_SIZE = 350; - const VirtualizedMessages: FC = ({ scrollRef, - virtualizerHandleRef, + virtualizerRef, + enabled, }) => { - const messageCount = useAuiState((s) => s.thread.messages.length); + // Two separate subscriptions so a prolonged streaming session doesn't force + // a re-render for every token on selectors we don't care about. + const messageIds = useAuiState((s) => s.thread.messages.map((m) => m.id)); + const messageCount = messageIds.length; + const isRunning = useAuiState((s) => s.thread.isRunning); + + // Pin the streaming assistant message into the rendered range even when the + // user has scrolled away. Without this, scrolling up during a stream would + // unmount the message — which stops useMessagePartText from receiving chunks + // and flips hover/last state. Only active while thread.isRunning so we don't + // permanently pin once the run finishes. + const pinnedIndex = isRunning && messageCount > 0 ? messageCount - 1 : -1; + const rangeExtractor = useCallback( + (range: Parameters[0]) => { + const base = defaultRangeExtractor(range); + if (pinnedIndex < 0 || base.includes(pinnedIndex)) return base; + // Insert while preserving ascending order so react-virtual renders them + // in DOM order (it otherwise re-orders absolutely-positioned rows + // naturally, but keeping the array sorted is cheaper for its internal + // diffing). + const next = [...base, pinnedIndex].sort((a, b) => a - b); + return next; + }, + [pinnedIndex], + ); - const virtualizer = useVirtualizer({ + const virtualizer = useVirtualizer({ count: messageCount, getScrollElement: () => scrollRef.current, estimateSize: () => ROW_ESTIMATE_SIZE, - overscan: 5, + overscan: OVERSCAN, + enabled, + // Stable keys prevent measurementsCache invalidation when branches flip or + // messages are inserted/reordered. + getItemKey: (index) => messageIds[index] ?? index, + // Leave ROW_ESTIMATE_SIZE-worth of breathing room at the top when we + // scrollToIndex(..., { align: 'start' }). This is what makes the new user + // message sit nicely under any header/padding instead of the viewport edge. + scrollPaddingStart: SCROLL_PADDING_START, + // Modern browsers expose a native `scrollend` event; opting into it lets + // TanStack skip its polling fallback for isScrolling. + useScrollendEvent: true, + rangeExtractor, }); - // Publish imperative handle for ThreadScrollController. We expose the - // virtualizer's own scrollToIndex so callers don't have to recompute offsets. + // Publish the instance up to Thread so ThreadScrollController can reach it + // without threading props through every render. useEffect(() => { - virtualizerHandleRef.current = { - scrollToIndex: (index, options) => virtualizer.scrollToIndex(index, options), - }; + virtualizerRef.current = virtualizer; return () => { - virtualizerHandleRef.current = null; + virtualizerRef.current = null; }; - }, [virtualizer, virtualizerHandleRef]); + }, [virtualizer, virtualizerRef]); if (messageCount === 0) return null; + // Docs pattern: translate the *wrapper* once by items[0].start, and render + // rows in natural document order inside it. That's cheaper than translating + // every row individually and plays nicely with contain:strict on the parent. + const items = virtualizer.getVirtualItems(); + return (
= ({ position: "relative", }} > - {virtualizer.getVirtualItems().map((virtualRow) => ( -
- -
- ))} +
+ {items.map((virtualRow) => ( +
+ +
+ ))} +
); }; From 7e8fcb963cbad12f5b6de7c1a125908c9912f0fc Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:21:12 +0000 Subject: [PATCH 04/12] refactor(thread): trim bloat from virtualized thread - Drop messages.map((m) => m.id) subscription; read ids lazily inside getItemKey. Previously every streaming token allocated a new array and re-rendered VirtualizedMessages even though the resulting ids were identical. Now we subscribe only to messages.length + isRunning (both primitives) so Object.is works. - Replace Parameters[0] with the exported Range type. - Collapse verbose JSDoc and inline comments that duplicated adjacent code. Net -35 lines. --- src/components/assistant-ui/thread.tsx | 133 +++++++++---------------- 1 file changed, 49 insertions(+), 84 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 7d60f549..bf85385e 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -43,6 +43,7 @@ import { import { defaultRangeExtractor, useVirtualizer, + type Range, type Virtualizer, } from "@tanstack/react-virtual"; @@ -120,21 +121,15 @@ interface ThreadProps { items?: Item[]; } -// Shared type alias: the @tanstack/react-virtual instance used throughout this -// file. We lift it to Thread so the scroll controller can call scrollToIndex -// directly without going through a handle wrapper. +// Shared type: the @tanstack/react-virtual instance. Lifted to Thread so the +// scroll controller can call scrollToIndex directly. type ThreadVirtualizer = Virtualizer; -// Estimated row height used before the virtualizer measures a row's actual -// size. Matches the previous value — tuned for a typical assistant turn with -// some markdown and a tool call. +// Default row height before the virtualizer measures actual sizes. const ROW_ESTIMATE_SIZE = 350; -// Breathing room above a user message when we scrollToIndex(..., { align: 'start' }). -// Without this the message is flush against the viewport's top edge, which -// looks wrong and hides MessageContextBadges on touch. +// Breathing room above a user message when scrollToIndex(align:'start'). const SCROLL_PADDING_START = 16; -// Rows kept mounted outside the visible viewport (above and below) so smooth -// scrolling doesn't flash blank regions. 5 matches TanStack's own docs example. +// Rows kept mounted outside the viewport for smooth scrolling. const OVERSCAN = 5; export const Thread: FC = ({ items = [] }) => { @@ -160,14 +155,9 @@ export const Thread: FC = ({ items = [] }) => { {/* Custom scroll container — intentionally NOT ThreadPrimitive.Viewport. assistant-ui's Viewport installs a ResizeObserver + subtree MutationObserver - on the scroll container and MessagePrimitive.Root injects inline min-height - on the last assistant message via ThreadViewportSlack. Both fight our - @tanstack/react-virtual virtualizer. We own scroll behavior here. - - - `contain: strict` lets the browser treat this subtree as an independent - layout/paint root, cutting reflow cost during streaming. - - `overflow-anchor: none` disables the browser's native scroll-anchoring, - which otherwise fights the virtualizer when row sizes change. + that fights @tanstack/react-virtual. We own scroll behavior here instead. + • contain:strict = independent layout/paint root, cuts reflow during streaming + • overflow-anchor:none = disables browser scroll anchoring (fights the virtualizer) */}
= ({ virtualizerRef, @@ -225,47 +208,33 @@ const ThreadScrollController: FC = ({ const scrollToLastIndex = useCallback( (align: "start" | "end", behavior: ScrollBehavior) => { + // Read live count from aui state; closure-based counts are unsafe + // against sibling-effect ordering (VirtualizedMessages may not have + // committed yet when the event fires). const count = aui?.thread()?.getState()?.messages.length ?? 0; if (count <= 0) return; - // Read the live count from aui state instead of any closure so we're - // safe against sibling effect ordering (VirtualizedMessages may not - // have committed the new count when this event fires). - const targetIndex = count - 1; - // rAF lets react-virtual see the new row count in its next commit - // before we scroll; scrollToIndex itself then retries internally for - // rows whose size hasn't been measured yet. + // rAF so react-virtual sees the new row count before we scroll; + // scrollToIndex retries internally for unmeasured rows. requestAnimationFrame(() => { - virtualizerRef.current?.scrollToIndex(targetIndex, { align, behavior }); + virtualizerRef.current?.scrollToIndex(count - 1, { align, behavior }); }); }, [aui, virtualizerRef], ); - useAuiEvent("thread.initialize", () => { - scrollToLastIndex("end", "instant"); - }); - - useAuiEvent("threadListItem.switchedTo", () => { - scrollToLastIndex("end", "instant"); - }); - - useAuiEvent("thread.runStart", () => { - // At runStart the just-sent user message is at messages.length - 1. - // Anchoring it at the top mirrors the top-turnAnchor feel we previously - // got from assistant-ui, without the min-height injection that was - // fighting our virtualizer. - scrollToLastIndex("start", "smooth"); - }); + useAuiEvent("thread.initialize", () => scrollToLastIndex("end", "instant")); + useAuiEvent("threadListItem.switchedTo", () => scrollToLastIndex("end", "instant")); + useAuiEvent("thread.runStart", () => scrollToLastIndex("start", "smooth")); return null; }; /** - * Replacement for the hover tracking that `MessagePrimitive.Root` normally wires - * up (via ThreadPrimitiveViewportSlack). Without this, ActionBarPrimitive - * autohide="not-last" and MessagePrimitive.If lastOrHover would be permanently - * hidden on non-last messages. We forward hover state into the aui message - * client directly, which is what MessageRoot does under the hood. + * Forwards hover state to the aui message client — the only side-effect of + * MessagePrimitive.Root we still need (powers ActionBarPrimitive autohide and + * MessagePrimitive.If lastOrHover). We dropped MessagePrimitive.Root itself + * because its ThreadViewportSlack wrapper injects inline min-height that the + * virtualizer's measureElement reads as the row size. */ const useMessageHoverHandlers = () => { const aui = useAui(); @@ -289,28 +258,25 @@ const VirtualizedMessages: FC = ({ virtualizerRef, enabled, }) => { - // Two separate subscriptions so a prolonged streaming session doesn't force - // a re-render for every token on selectors we don't care about. - const messageIds = useAuiState((s) => s.thread.messages.map((m) => m.id)); - const messageCount = messageIds.length; + const aui = useAui(); + // Subscribing to length + isRunning (both primitives) is enough to drive + // re-renders. getItemKey reads ids lazily from live aui state, so we don't + // need to subscribe to the whole messages array (which would allocate a + // new array identity on every streaming token). + const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); // Pin the streaming assistant message into the rendered range even when the // user has scrolled away. Without this, scrolling up during a stream would - // unmount the message — which stops useMessagePartText from receiving chunks - // and flips hover/last state. Only active while thread.isRunning so we don't - // permanently pin once the run finishes. + // unmount the message — stopping useMessagePartText from receiving chunks + // and flipping hover/last state. Only active while thread.isRunning. const pinnedIndex = isRunning && messageCount > 0 ? messageCount - 1 : -1; const rangeExtractor = useCallback( - (range: Parameters[0]) => { + (range: Range) => { const base = defaultRangeExtractor(range); if (pinnedIndex < 0 || base.includes(pinnedIndex)) return base; - // Insert while preserving ascending order so react-virtual renders them - // in DOM order (it otherwise re-orders absolutely-positioned rows - // naturally, but keeping the array sorted is cheaper for its internal - // diffing). - const next = [...base, pinnedIndex].sort((a, b) => a - b); - return next; + // Keep ascending order so react-virtual's internal diff is linear. + return [...base, pinnedIndex].sort((a, b) => a - b); }, [pinnedIndex], ); @@ -321,15 +287,14 @@ const VirtualizedMessages: FC = ({ estimateSize: () => ROW_ESTIMATE_SIZE, overscan: OVERSCAN, enabled, - // Stable keys prevent measurementsCache invalidation when branches flip or - // messages are inserted/reordered. - getItemKey: (index) => messageIds[index] ?? index, - // Leave ROW_ESTIMATE_SIZE-worth of breathing room at the top when we - // scrollToIndex(..., { align: 'start' }). This is what makes the new user - // message sit nicely under any header/padding instead of the viewport edge. + // Stable keys by aui message id prevent measurementsCache invalidation + // on branch flips / reorders. Read live from aui state so this selector + // doesn't need to subscribe to the full messages array. + getItemKey: (index) => + aui?.thread()?.getState()?.messages[index]?.id ?? index, + // Breathing room above a user message when scrollToIndex(align:'start'). scrollPaddingStart: SCROLL_PADDING_START, - // Modern browsers expose a native `scrollend` event; opting into it lets - // TanStack skip its polling fallback for isScrolling. + // Native scrollend event on modern browsers skips TanStack's polling. useScrollendEvent: true, rangeExtractor, }); @@ -345,9 +310,9 @@ const VirtualizedMessages: FC = ({ if (messageCount === 0) return null; - // Docs pattern: translate the *wrapper* once by items[0].start, and render - // rows in natural document order inside it. That's cheaper than translating - // every row individually and plays nicely with contain:strict on the parent. + // Docs pattern: translate the wrapper once by items[0].start, render rows + // in natural document order inside. Cheaper than per-row transforms, works + // well with contain:strict on the parent. const items = virtualizer.getVirtualItems(); return ( From 96dc7839f4ffe9de752424a81890e99cc032caff Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:24:01 +0000 Subject: [PATCH 05/12] fix(thread): per-row translateY for non-contiguous virtual items The TanStack "dynamic" docs example uses a single wrapper translate (translateY(items[0].start)) and renders rows in natural document order inside. That optimization is only correct when items is a contiguous index range. Our rangeExtractor pins the streaming last assistant message while thread.isRunning, so items can be non-contiguous: e.g. [5,6,7,8,9,10,42] when the user has scrolled up during a stream. Under the wrapper-only translate, row 42 would render right after row 10's natural flow at the wrong y-offset. Revert to per-row absolute positioning (each row translated by its own virtualRow.start), which stays correct for any rangeExtractor output. Minor extra transform cost, but rangeExtractor-pinning was the whole point of this refactor. Caught by cubic-dev-ai reviewing PR #394. --- src/components/assistant-ui/thread.tsx | 50 +++++++++++++------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index bf85385e..97c74895 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -310,9 +310,12 @@ const VirtualizedMessages: FC = ({ if (messageCount === 0) return null; - // Docs pattern: translate the wrapper once by items[0].start, render rows - // in natural document order inside. Cheaper than per-row transforms, works - // well with contain:strict on the parent. + // Per-row absolute positioning (not the single-wrapper translate shown in + // TanStack's "dynamic" docs example). The wrapper-translate pattern assumes + // items is a CONTIGUOUS index range, but our rangeExtractor pins the + // streaming last assistant message so items can look like [5,6,7,...,42] + // while the user scrolls up. Translating per-row keeps each row at its true + // virtualRow.start regardless of gaps. const items = virtualizer.getVirtualItems(); return ( @@ -323,28 +326,25 @@ const VirtualizedMessages: FC = ({ position: "relative", }} > -
- {items.map((virtualRow) => ( -
- -
- ))} -
+ {items.map((virtualRow) => ( +
+ +
+ ))}
); }; From f704e527d66a6ed6ebf64e8d2d923b831f979812 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 18 Apr 2026 00:01:00 +0000 Subject: [PATCH 06/12] fix(thread): bottom slack + reliable thread-switch scroll Three issues from user testing the virtualized thread: 1. Sending a new message didn't anchor it at the top of the viewport when there wasn't enough content below. virtualizer.scrollToIndex clamps to getMaxScrollOffset, which is tight against the last row on a fresh send. Added paddingEnd: 800 so there's extra scrollable space below the last message for the align:'start' scroll to land properly (this is the Tanstack-native way to get the 'slack' that assistant-ui's ViewportSlack was trying to add via min-height injection). 2. Switching threads didn't scroll to bottom. The threadListItem .switchedTo event fires BEFORE the new thread's messages are imported (it's triggered by the thread id changing on the list item), so reading messages.length inside the handler returned 0 or the old count. Capture the threadId in state via the event, then react in a useEffect after React commits the new messages array. Also note thread.initialize is per-runtime-instance and fires only once per thread, so it can't cover re-switches to already-loaded threads \u2014 hence the split. 3. Bumped the rAF dance to two frames so measureElement has a chance to register sizes before scrollToIndex reads them. scrollToIndex retries internally for unmeasured rows anyway, but two frames lets us land on actual measured offsets on the first attempt more often. Smooth scrolling on runStart already works via { behavior: 'smooth' } \u2014 Tanstack's scrollToIndex uses native CSS smooth scrolling by default. If we want a custom ease (duration-controlled) we'd need a scrollToFn override, which we're not doing. --- src/components/assistant-ui/thread.tsx | 47 ++++++++++++++++++++++---- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 97c74895..822d49f5 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -129,6 +129,12 @@ type ThreadVirtualizer = Virtualizer; const ROW_ESTIMATE_SIZE = 350; // Breathing room above a user message when scrollToIndex(align:'start'). const SCROLL_PADDING_START = 16; +// Bottom slack — extra space appended to the virtualized region so a freshly +// sent user message can be anchored at the top of the viewport even when +// the assistant hasn't streamed any content below it yet. Without this, +// scrollToIndex(lastIndex, { align: 'start' }) clamps because the last +// message is also the bottom of the scrollable area. +const BOTTOM_SLACK_PX = 800; // Rows kept mounted outside the viewport for smooth scrolling. const OVERSCAN = 5; @@ -196,8 +202,9 @@ interface ThreadScrollControllerProps { /** * Owns scroll behavior for the virtualized thread. - * • thread.initialize → bottom, instant - * • threadListItem.switchedTo → bottom, instant + * • thread.initialize → bottom, instant (first hydration of a thread runtime) + * • thread id change → bottom, instant (covers switching to a previously + * loaded thread, where thread.initialize won't fire again) * • thread.runStart → anchor the new user message at the top, smooth * • Streaming: user owns scroll */ @@ -205,25 +212,46 @@ const ThreadScrollController: FC = ({ virtualizerRef, }) => { const aui = useAui(); + // Track the current thread id via the threadListItem.switchedTo event + // payload. We can't read it from thread state (no id field there) and we + // can't scroll inside the event itself because the new thread's messages + // haven't been imported yet \u2014 so we stash the id and react in an effect. + const [activeThreadId, setActiveThreadId] = useState(null); const scrollToLastIndex = useCallback( (align: "start" | "end", behavior: ScrollBehavior) => { // Read live count from aui state; closure-based counts are unsafe // against sibling-effect ordering (VirtualizedMessages may not have - // committed yet when the event fires). + // committed yet when this fires). const count = aui?.thread()?.getState()?.messages.length ?? 0; if (count <= 0) return; - // rAF so react-virtual sees the new row count before we scroll; - // scrollToIndex retries internally for unmeasured rows. + // Two RAFs: first lets react-virtual commit the new row count, second + // gives measureElement a chance to register row sizes before we scroll. + // scrollToIndex itself also retries internally for unmeasured rows. requestAnimationFrame(() => { - virtualizerRef.current?.scrollToIndex(count - 1, { align, behavior }); + requestAnimationFrame(() => { + virtualizerRef.current?.scrollToIndex(count - 1, { align, behavior }); + }); }); }, [aui, virtualizerRef], ); + // First hydration of a thread runtime (per-instance event, fires once). useAuiEvent("thread.initialize", () => scrollToLastIndex("end", "instant")); - useAuiEvent("threadListItem.switchedTo", () => scrollToLastIndex("end", "instant")); + + // Switching to a previously loaded thread does NOT re-fire thread.initialize + // (each thread runtime caches _isInitialized). Capture the threadId here, then + // react in an effect after React commits the new messages array. + useAuiEvent("threadListItem.switchedTo", ({ threadId }) => { + setActiveThreadId(threadId); + }); + useEffect(() => { + if (!activeThreadId) return; + scrollToLastIndex("end", "instant"); + }, [activeThreadId, scrollToLastIndex]); + + // New user message sent \u2192 anchor at top with smooth animation. useAuiEvent("thread.runStart", () => scrollToLastIndex("start", "smooth")); return null; @@ -294,6 +322,11 @@ const VirtualizedMessages: FC = ({ aui?.thread()?.getState()?.messages[index]?.id ?? index, // Breathing room above a user message when scrollToIndex(align:'start'). scrollPaddingStart: SCROLL_PADDING_START, + // Bottom slack — padding appended after the last row so scrollToIndex( + // lastIndex, { align: 'start' }) can land the new user message at the top + // even when the assistant hasn't streamed content below it yet. Without + // this padding, getMaxScrollOffset clamps the scroll mid-viewport. + paddingEnd: BOTTOM_SLACK_PX, // Native scrollend event on modern browsers skips TanStack's polling. useScrollendEvent: true, rangeExtractor, From 8553adef357cf7547c0a746487422c03192ba979 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 18 Apr 2026 02:10:50 +0000 Subject: [PATCH 07/12] refactor(thread): drive thread-switch scroll via useAuiState subscription Replace the threadListItem.switchedTo -> useState -> useEffect chain with a direct useAuiState subscription on (threadListItem.id, thread.messages.length). Rationale, cross-referenced against assistant-ui source: - The upstream useThreadViewportAutoScroll listens to threadListItem.switchedTo but scrolls el.scrollHeight, not an index. That works for browser-native scroll but misbehaves with the virtualizer's measured ledger. - thread.initialize fires BEFORE repository.import(), so reading messages.length inside the event always sees 0. Subscribing to messages.length via useAuiState lets the effect re-run when the count actually lands. - threadListItem.id is the canonical thread identifier on the store proxy (confirmed via scope-registration.d.ts / thread-list-item.d.ts in @assistant-ui/core). thread.threadId only exists on the core runtime's ThreadState, not the store-projected one. A ref-backed settled-thread guard ensures A -> B -> A re-scrolls on the second visit, without yanking scroll on every message tick during streaming. runStart remains event-driven for smooth top-anchor. --- src/components/assistant-ui/thread.tsx | 90 ++++++++++++++++---------- 1 file changed, 56 insertions(+), 34 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 822d49f5..40bde797 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -201,29 +201,40 @@ interface ThreadScrollControllerProps { } /** - * Owns scroll behavior for the virtualized thread. - * • thread.initialize → bottom, instant (first hydration of a thread runtime) - * • thread id change → bottom, instant (covers switching to a previously - * loaded thread, where thread.initialize won't fire again) - * • thread.runStart → anchor the new user message at the top, smooth - * • Streaming: user owns scroll + * Owns scroll behavior for the virtualized thread. Two concerns: + * + * 1. LOAD / THREAD-SWITCH -> bottom, instant. Tracked via `threadListItem.id` + * + `thread.messages.length`. A ref-backed "last scrolled for" guard + * ensures we only snap to bottom on thread identity changes — not on every + * message count tick. Reading `messageCount` as a dep is what lets us + * catch the initial hydration: the runtime emits `thread.initialize` + * BEFORE `repository.import()`, so at event time count is 0. By making + * the effect depend on count, it re-runs once messages land and we can + * commit the scroll. + * + * Mirrors assistant-ui's upstream `useThreadViewportAutoScroll`, except + * we target the virtualizer's measurement ledger instead of the raw + * `el.scrollHeight`. + * + * 2. NEW USER MESSAGE (thread.runStart) -> anchor at top, smooth. Event-driven + * because we need the transition, not the steady state. Reads the latest + * `messageCount` from the subscription via `useEffectEvent` semantics of + * `useAuiEvent`. Organic growth during streaming is deliberately ignored — + * the user owns scroll once anchored. */ const ThreadScrollController: FC = ({ virtualizerRef, }) => { - const aui = useAui(); - // Track the current thread id via the threadListItem.switchedTo event - // payload. We can't read it from thread state (no id field there) and we - // can't scroll inside the event itself because the new thread's messages - // haven't been imported yet \u2014 so we stash the id and react in an effect. - const [activeThreadId, setActiveThreadId] = useState(null); - - const scrollToLastIndex = useCallback( - (align: "start" | "end", behavior: ScrollBehavior) => { - // Read live count from aui state; closure-based counts are unsafe - // against sibling-effect ordering (VirtualizedMessages may not have - // committed yet when this fires). - const count = aui?.thread()?.getState()?.messages.length ?? 0; + // Reactive subscriptions: both drive the load/switch effect below. Using + // useAuiState over useAuiEvent lets us observe the settled state rather + // than racing event emission against React commits. `threadListItem.id` + // is the canonical thread identifier on the store proxy; `thread.threadId` + // only exists on the core runtime's ThreadState and isn't exposed here. + const threadId = useAuiState((s) => s.threadListItem.id); + const messageCount = useAuiState((s) => s.thread.messages.length); + + const scrollToIndex = useCallback( + (count: number, align: "start" | "end", behavior: ScrollBehavior) => { if (count <= 0) return; // Two RAFs: first lets react-virtual commit the new row count, second // gives measureElement a chance to register row sizes before we scroll. @@ -234,25 +245,36 @@ const ThreadScrollController: FC = ({ }); }); }, - [aui, virtualizerRef], + [virtualizerRef], ); - // First hydration of a thread runtime (per-instance event, fires once). - useAuiEvent("thread.initialize", () => scrollToLastIndex("end", "instant")); + // Track the thread we're currently "settled" on. Reset whenever the id + // changes so that visiting A → B → A re-scrolls A's bottom. + const settledThreadIdRef = useRef(null); - // Switching to a previously loaded thread does NOT re-fire thread.initialize - // (each thread runtime caches _isInitialized). Capture the threadId here, then - // react in an effect after React commits the new messages array. - useAuiEvent("threadListItem.switchedTo", ({ threadId }) => { - setActiveThreadId(threadId); - }); useEffect(() => { - if (!activeThreadId) return; - scrollToLastIndex("end", "instant"); - }, [activeThreadId, scrollToLastIndex]); - - // New user message sent \u2192 anchor at top with smooth animation. - useAuiEvent("thread.runStart", () => scrollToLastIndex("start", "smooth")); + // Reset settle marker whenever the visible thread id changes; this makes + // A → B → A trigger a bottom-scroll on the second visit to A. + if (settledThreadIdRef.current !== threadId) { + settledThreadIdRef.current = null; + } + // Wait for messages to land before committing the "settled" marker — + // ensureInitialized() fires before repository.import(), so on first + // mount messageCount is 0 and flips to N on the next commit. + if (messageCount <= 0) return; + if (settledThreadIdRef.current === threadId) return; + settledThreadIdRef.current = threadId; + scrollToIndex(messageCount, "end", "instant"); + }, [threadId, messageCount, scrollToIndex]); + + // New user message sent -> anchor at top with smooth animation. Event-driven + // because we need to react on the transition, not the steady state. The + // handler reads `messageCount` via useEffectEvent (always latest), and + // since runStart fires after the user message is appended to the store, + // `messageCount` already reflects the new row. + useAuiEvent("thread.runStart", () => { + scrollToIndex(messageCount, "start", "smooth"); + }); return null; }; From 633c181838c98a741971c1c6bd9ba53c0eedda59 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 18 Apr 2026 02:23:12 +0000 Subject: [PATCH 08/12] fix(thread): read live message count in runStart handler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Capy review flagged a timing bug: `thread.runStart` is emitted synchronously inside the runtime's append -> startRun path, between the repository mutation that appends the user message and React's commit phase. Because `useEffectEvent` (the underlying mechanism of `useAuiEvent`) updates its ref in `useInsertionEffect` (commit phase), at runStart time the ref still points to the previous closure. That closure captured the pre-send `messageCount` from the `useAuiState` subscription, so scrolling to `messageCount` anchored the *previous* assistant message rather than the newly-sent user message. Fix: read the count imperatively from the store inside the handler via `aui.thread().getState().messages.length`. The store is always current at event emission time — only the React subscriptions are one commit behind. The subscription-driven load/switch effect still uses the reactive `messageCount` because it runs post-commit where closure staleness doesn't apply. Verified against @assistant-ui/core@0.1.14: - local-thread-runtime-core.ts startRun() emits runStart directly after repository.addOrUpdateMessage with no intervening _notifySubscribers tick. - use-effect-event@2.0.3 updates its ref inside useInsertionEffect. --- src/components/assistant-ui/thread.tsx | 29 +++++++++++++++++--------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 40bde797..8646db1d 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -216,15 +216,18 @@ interface ThreadScrollControllerProps { * we target the virtualizer's measurement ledger instead of the raw * `el.scrollHeight`. * - * 2. NEW USER MESSAGE (thread.runStart) -> anchor at top, smooth. Event-driven - * because we need the transition, not the steady state. Reads the latest - * `messageCount` from the subscription via `useEffectEvent` semantics of - * `useAuiEvent`. Organic growth during streaming is deliberately ignored — - * the user owns scroll once anchored. + * 2. NEW USER MESSAGE (thread.runStart) -> anchor at top, smooth. Event-driven. + * The count is read imperatively from the store inside the handler, NOT + * from the `messageCount` subscription closure: runStart fires + * synchronously between the repository mutation and React's commit, so + * `useEffectEvent`'s ref still holds the stale closure with the old count. + * Organic growth during streaming is deliberately ignored — the user owns + * scroll once anchored. */ const ThreadScrollController: FC = ({ virtualizerRef, }) => { + const aui = useAui(); // Reactive subscriptions: both drive the load/switch effect below. Using // useAuiState over useAuiEvent lets us observe the settled state rather // than racing event emission against React commits. `threadListItem.id` @@ -268,12 +271,18 @@ const ThreadScrollController: FC = ({ }, [threadId, messageCount, scrollToIndex]); // New user message sent -> anchor at top with smooth animation. Event-driven - // because we need to react on the transition, not the steady state. The - // handler reads `messageCount` via useEffectEvent (always latest), and - // since runStart fires after the user message is appended to the store, - // `messageCount` already reflects the new row. + // because we need to react on the transition, not the steady state. We must + // read the count *imperatively from the store* here, not from the + // `messageCount` closure: `runStart` is emitted synchronously inside the + // runtime's append->startRun path, AFTER `repository.addOrUpdateMessage` + // appends the user message but BEFORE React commits the re-render. At that + // moment `useEffectEvent`'s ref (updated in useInsertionEffect, commit-phase) + // still points at the previous closure with the old count, so using the + // subscription value here would scroll to the last *previous* message. + // Reading `aui.thread().getState().messages.length` bypasses that window. useAuiEvent("thread.runStart", () => { - scrollToIndex(messageCount, "start", "smooth"); + const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; + scrollToIndex(liveCount, "start", "smooth"); }); return null; From 2f5a44ddadc4569d3a8aba678977656e0699c4b2 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 18 Apr 2026 03:00:00 +0000 Subject: [PATCH 09/12] refactor(thread): swap @tanstack/react-virtual for virtua MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Virtua is a chat-native virtualizer: it ships `scrollToIndex({ smooth, align })`, `keepMounted` (the streaming-pin we built manually), `viewportSize`, and absolute-positioning of items internally. This deletes most of our TanStack Virtual compensation code. Changes: - package.json: drop @tanstack/react-virtual@^3.13.23, add virtua@^0.49.1. - thread.tsx: net -92 lines in this file. - Delete ROW_ESTIMATE_SIZE, SCROLL_PADDING_START, BOTTOM_SLACK_PX, OVERSCAN constants. Virtua auto-estimates, handles scroll alignment natively, and replaces BOTTOM_SLACK_PX with the per-row minHeight trick from virtua's canonical chatbot demo. - Delete `rangeExtractor` + `defaultRangeExtractor` + `Range` import. Replaced by `keepMounted={[streamingAssistantIndex]}` prop. - Delete the `mounted` gate + `enabled` prop on VirtualizedMessages. Virtua handles SSR/null-scrollRef lazily. - Delete per-row absolute positioning + manual translateY + measureElement wiring. Virtua positions items internally. - Delete scrollPaddingTop CSS on the scroll container. Virtua's scrollToIndex(align:'start') lands rows flush without needing it. - Replace manual paddingEnd bottom-slack with the virtua-demo pattern: capture viewportSize on thread.runStart, measure the fresh user message's height via ResizeObserver, apply minHeight = viewportSize - userHeight to the streaming assistant row. Cleans up when isRunning flips false. Keeps the assistant-ui mediation bits (custom scroll container bypassing ThreadPrimitive.Viewport, useMessageHoverHandlers, ThreadPrimitive.MessageByIndex components map) — those are orthogonal to the virtualizer choice. Reference: https://github.com/inokawa/virtua Chat.stories.tsx --- package.json | 2 +- src/components/assistant-ui/thread.tsx | 290 +++++++++---------------- 2 files changed, 100 insertions(+), 192 deletions(-) diff --git a/package.json b/package.json index d2821d64..d69bf3ce 100644 --- a/package.json +++ b/package.json @@ -93,7 +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", + "virtua": "^0.49.1", "@tiptap/core": "3.22.3", "@tiptap/extension-code-block": "3.22.3", "@tiptap/extension-highlight": "3.22.3", diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 8646db1d..9fc4d660 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -40,14 +40,9 @@ import { useMessagePartText, useAuiState, } from "@assistant-ui/react"; -import { - defaultRangeExtractor, - useVirtualizer, - type Range, - type Virtualizer, -} from "@tanstack/react-virtual"; +import { Virtualizer, type VirtualizerHandle } from "virtua"; -import type { FC, MutableRefObject, RefObject } from "react"; +import type { FC, RefObject } from "react"; import { createContext, useContext } from "react"; import { useEffect, useRef, useState, useMemo, useCallback } from "react"; @@ -121,35 +116,12 @@ interface ThreadProps { items?: Item[]; } -// Shared type: the @tanstack/react-virtual instance. Lifted to Thread so the -// scroll controller can call scrollToIndex directly. -type ThreadVirtualizer = Virtualizer; - -// Default row height before the virtualizer measures actual sizes. -const ROW_ESTIMATE_SIZE = 350; -// Breathing room above a user message when scrollToIndex(align:'start'). -const SCROLL_PADDING_START = 16; -// Bottom slack — extra space appended to the virtualized region so a freshly -// sent user message can be anchored at the top of the viewport even when -// the assistant hasn't streamed any content below it yet. Without this, -// scrollToIndex(lastIndex, { align: 'start' }) clamps because the last -// message is also the bottom of the scrollable area. -const BOTTOM_SLACK_PX = 800; -// Rows kept mounted outside the viewport for smooth scrolling. -const OVERSCAN = 5; +type ThreadVirtualizer = VirtualizerHandle; export const Thread: FC = ({ items = [] }) => { const viewportRef = useRef(null); - // Lifted so ThreadScrollController can imperatively scroll to an index - // without a separate "handle" abstraction. const virtualizerRef = useRef(null); - - // Gate useVirtualizer behind mount so SSR/first-paint doesn't compute a - // range against a null scroll element (TanStack's `enabled` option). - const [mounted, setMounted] = useState(false); - useEffect(() => { - setMounted(true); - }, []); + const captureBlankRef = useRef<(() => void) | null>(null); return ( = ({ items = [] }) => { > {/* Custom scroll container — intentionally NOT ThreadPrimitive.Viewport. - assistant-ui's Viewport installs a ResizeObserver + subtree MutationObserver - that fights @tanstack/react-virtual. We own scroll behavior here instead. - • contain:strict = independent layout/paint root, cuts reflow during streaming - • overflow-anchor:none = disables browser scroll anchoring (fights the virtualizer) + assistant-ui's Viewport installs a ResizeObserver + subtree + MutationObserver that fights any virtualizer. We own scroll behavior. */}
thread.isLoading}> @@ -184,9 +150,12 @@ export const Thread: FC = ({ items = [] }) => { + -
@@ -197,104 +166,49 @@ export const Thread: FC = ({ items = [] }) => { }; interface ThreadScrollControllerProps { - virtualizerRef: MutableRefObject; + virtualizerRef: RefObject; + captureBlankRef: RefObject<(() => void) | null>; } -/** - * Owns scroll behavior for the virtualized thread. Two concerns: - * - * 1. LOAD / THREAD-SWITCH -> bottom, instant. Tracked via `threadListItem.id` - * + `thread.messages.length`. A ref-backed "last scrolled for" guard - * ensures we only snap to bottom on thread identity changes — not on every - * message count tick. Reading `messageCount` as a dep is what lets us - * catch the initial hydration: the runtime emits `thread.initialize` - * BEFORE `repository.import()`, so at event time count is 0. By making - * the effect depend on count, it re-runs once messages land and we can - * commit the scroll. - * - * Mirrors assistant-ui's upstream `useThreadViewportAutoScroll`, except - * we target the virtualizer's measurement ledger instead of the raw - * `el.scrollHeight`. - * - * 2. NEW USER MESSAGE (thread.runStart) -> anchor at top, smooth. Event-driven. - * The count is read imperatively from the store inside the handler, NOT - * from the `messageCount` subscription closure: runStart fires - * synchronously between the repository mutation and React's commit, so - * `useEffectEvent`'s ref still holds the stale closure with the old count. - * Organic growth during streaming is deliberately ignored — the user owns - * scroll once anchored. - */ const ThreadScrollController: FC = ({ virtualizerRef, + captureBlankRef, }) => { const aui = useAui(); - // Reactive subscriptions: both drive the load/switch effect below. Using - // useAuiState over useAuiEvent lets us observe the settled state rather - // than racing event emission against React commits. `threadListItem.id` - // is the canonical thread identifier on the store proxy; `thread.threadId` - // only exists on the core runtime's ThreadState and isn't exposed here. const threadId = useAuiState((s) => s.threadListItem.id); const messageCount = useAuiState((s) => s.thread.messages.length); - const scrollToIndex = useCallback( - (count: number, align: "start" | "end", behavior: ScrollBehavior) => { + const scrollToLastIndex = useCallback( + (count: number, align: "start" | "end", smooth: boolean) => { if (count <= 0) return; - // Two RAFs: first lets react-virtual commit the new row count, second - // gives measureElement a chance to register row sizes before we scroll. - // scrollToIndex itself also retries internally for unmeasured rows. requestAnimationFrame(() => { - requestAnimationFrame(() => { - virtualizerRef.current?.scrollToIndex(count - 1, { align, behavior }); - }); + virtualizerRef.current?.scrollToIndex(count - 1, { align, smooth }); }); }, [virtualizerRef], ); - // Track the thread we're currently "settled" on. Reset whenever the id - // changes so that visiting A → B → A re-scrolls A's bottom. const settledThreadIdRef = useRef(null); useEffect(() => { - // Reset settle marker whenever the visible thread id changes; this makes - // A → B → A trigger a bottom-scroll on the second visit to A. if (settledThreadIdRef.current !== threadId) { settledThreadIdRef.current = null; } - // Wait for messages to land before committing the "settled" marker — - // ensureInitialized() fires before repository.import(), so on first - // mount messageCount is 0 and flips to N on the next commit. if (messageCount <= 0) return; if (settledThreadIdRef.current === threadId) return; settledThreadIdRef.current = threadId; - scrollToIndex(messageCount, "end", "instant"); - }, [threadId, messageCount, scrollToIndex]); - - // New user message sent -> anchor at top with smooth animation. Event-driven - // because we need to react on the transition, not the steady state. We must - // read the count *imperatively from the store* here, not from the - // `messageCount` closure: `runStart` is emitted synchronously inside the - // runtime's append->startRun path, AFTER `repository.addOrUpdateMessage` - // appends the user message but BEFORE React commits the re-render. At that - // moment `useEffectEvent`'s ref (updated in useInsertionEffect, commit-phase) - // still points at the previous closure with the old count, so using the - // subscription value here would scroll to the last *previous* message. - // Reading `aui.thread().getState().messages.length` bypasses that window. + scrollToLastIndex(messageCount, "end", false); + }, [threadId, messageCount, scrollToLastIndex]); + useAuiEvent("thread.runStart", () => { + captureBlankRef.current?.(); const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; - scrollToIndex(liveCount, "start", "smooth"); + scrollToLastIndex(liveCount, "start", true); }); return null; }; -/** - * Forwards hover state to the aui message client — the only side-effect of - * MessagePrimitive.Root we still need (powers ActionBarPrimitive autohide and - * MessagePrimitive.If lastOrHover). We dropped MessagePrimitive.Root itself - * because its ThreadViewportSlack wrapper injects inline min-height that the - * virtualizer's measureElement reads as the row size. - */ const useMessageHoverHandlers = () => { const aui = useAui(); const onMouseEnter = useCallback(() => { @@ -308,108 +222,102 @@ const useMessageHoverHandlers = () => { interface VirtualizedMessagesProps { scrollRef: RefObject; - virtualizerRef: MutableRefObject; - enabled: boolean; + virtualizerRef: RefObject; + captureBlankRef: RefObject<(() => void) | null>; } const VirtualizedMessages: FC = ({ scrollRef, virtualizerRef, - enabled, + captureBlankRef, }) => { const aui = useAui(); - // Subscribing to length + isRunning (both primitives) is enough to drive - // re-renders. getItemKey reads ids lazily from live aui state, so we don't - // need to subscribe to the whole messages array (which would allocate a - // new array identity on every streaming token). const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); + const [blankSize, setBlankSize] = useState(0); + const [lastUserSize, setLastUserSize] = useState(0); + const [measuredUserElement, setMeasuredUserElement] = + useState(null); - // Pin the streaming assistant message into the rendered range even when the - // user has scrolled away. Without this, scrolling up during a stream would - // unmount the message — stopping useMessagePartText from receiving chunks - // and flipping hover/last state. Only active while thread.isRunning. - const pinnedIndex = isRunning && messageCount > 0 ? messageCount - 1 : -1; - const rangeExtractor = useCallback( - (range: Range) => { - const base = defaultRangeExtractor(range); - if (pinnedIndex < 0 || base.includes(pinnedIndex)) return base; - // Keep ascending order so react-virtual's internal diff is linear. - return [...base, pinnedIndex].sort((a, b) => a - b); - }, - [pinnedIndex], - ); + useEffect(() => { + captureBlankRef.current = () => { + setBlankSize(virtualizerRef.current?.viewportSize ?? 0); + }; + return () => { + captureBlankRef.current = null; + }; + }, [captureBlankRef, virtualizerRef]); - const virtualizer = useVirtualizer({ - count: messageCount, - getScrollElement: () => scrollRef.current, - estimateSize: () => ROW_ESTIMATE_SIZE, - overscan: OVERSCAN, - enabled, - // Stable keys by aui message id prevent measurementsCache invalidation - // on branch flips / reorders. Read live from aui state so this selector - // doesn't need to subscribe to the full messages array. - getItemKey: (index) => - aui?.thread()?.getState()?.messages[index]?.id ?? index, - // Breathing room above a user message when scrollToIndex(align:'start'). - scrollPaddingStart: SCROLL_PADDING_START, - // Bottom slack — padding appended after the last row so scrollToIndex( - // lastIndex, { align: 'start' }) can land the new user message at the top - // even when the assistant hasn't streamed content below it yet. Without - // this padding, getMaxScrollOffset clamps the scroll mid-viewport. - paddingEnd: BOTTOM_SLACK_PX, - // Native scrollend event on modern browsers skips TanStack's polling. - useScrollendEvent: true, - rangeExtractor, - }); + useEffect(() => { + if (!isRunning && blankSize !== 0) { + setBlankSize(0); + } + }, [isRunning, blankSize]); + + useEffect(() => { + if (!isRunning && lastUserSize !== 0) { + setLastUserSize(0); + } + }, [isRunning, lastUserSize]); - // Publish the instance up to Thread so ThreadScrollController can reach it - // without threading props through every render. useEffect(() => { - virtualizerRef.current = virtualizer; + if (!measuredUserElement) return; + const update = () => { + setLastUserSize(measuredUserElement.getBoundingClientRect().height); + }; + update(); + const resizeObserver = new ResizeObserver(update); + resizeObserver.observe(measuredUserElement); return () => { - virtualizerRef.current = null; + resizeObserver.disconnect(); }; - }, [virtualizer, virtualizerRef]); + }, [measuredUserElement]); - if (messageCount === 0) return null; + const measureUserMessageRef = useCallback((el: HTMLDivElement | null) => { + setMeasuredUserElement(el); + }, []); + + const streamingAssistantIndex = isRunning ? messageCount - 1 : -1; + const measuredUserIndex = isRunning ? messageCount - 2 : -1; + const assistantMinHeight = + streamingAssistantIndex >= 0 ? Math.max(0, blankSize - lastUserSize) : 0; - // Per-row absolute positioning (not the single-wrapper translate shown in - // TanStack's "dynamic" docs example). The wrapper-translate pattern assumes - // items is a CONTIGUOUS index range, but our rangeExtractor pins the - // streaming last assistant message so items can look like [5,6,7,...,42] - // while the user scrolls up. Translating per-row keeps each row at its true - // virtualRow.start regardless of gaps. - const items = virtualizer.getVirtualItems(); + if (messageCount === 0) return null; return ( -
= 0 ? [streamingAssistantIndex] : undefined + } > - {items.map((virtualRow) => ( -
- -
- ))} -
+ {Array.from({ length: messageCount }, (_, index) => { + const message = aui?.thread()?.getState()?.messages[index]; + const id = message?.id ?? index; + const role = message?.role; + const isStreamingAssistant = + role === "assistant" && index === streamingAssistantIndex; + const isMeasuredUser = role === "user" && index === measuredUserIndex; + + return ( +
+ +
+ ); + })} + ); }; From 197db373f50fc2734708020aa61b088fb933f804 Mon Sep 17 00:00:00 2001 From: urjitc <135136842+urjitc@users.noreply.github.com> Date: Sat, 18 Apr 2026 03:11:16 +0000 Subject: [PATCH 10/12] fix(thread): post-commit scroll via isRunning effect instead of runStart event MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Problem: user-message-to-top was intermittently broken. Two causes, both flagged by capy review: 1. `thread.runStart` fires synchronously inside the runtime's repository mutation, BEFORE React commits the render that mounts the new rows. One rAF is not enough lead time; on fast streams `scrollToIndex` could run against a stale row count. 2. `VirtualizedMessages` returned null when messageCount === 0, so on the first message of a brand-new thread the Virtualizer wasn't mounted yet when runStart fired. `virtualizerRef.current` was null and `viewportSize` read 0, so blank-slack was 0 and the user message couldn't actually reach the top of the viewport. Fix: - Drop the runStart event handler entirely. Replace with a useEffect keyed on `isRunning`: when it transitions false -> true we capture viewportSize, read the live message count, and call scrollToIndex inside rAF. `hasScrolledForRunRef` dedupes against streaming-token count ticks. React guarantees the effect runs after commit, so the Virtualizer is mounted and the row is in the DOM. - Remove the `if (messageCount === 0) return null` guard in VirtualizedMessages. Always mount the Virtualizer so its handle is available on the first message. Empty-state UI is already handled by sibling AuiIf wrappers in Thread (ThreadWelcome / skeleton). - Drop the now-unused useAuiEvent import. This follows the virtua chatbot demo pattern (setState -> rAF -> scrollToIndex) correctly for the first time — previously the event- driven path was firing one commit too early. --- src/components/assistant-ui/thread.tsx | 45 ++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 9fc4d660..0a66198d 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -35,7 +35,6 @@ import { MessagePrimitive, ThreadPrimitive, useAui, - useAuiEvent, useMessage, useMessagePartText, useAuiState, @@ -170,6 +169,25 @@ interface ThreadScrollControllerProps { captureBlankRef: RefObject<(() => void) | null>; } +/** + * Owns scroll behavior for the virtualized thread. Three concerns: + * + * 1. LOAD / THREAD-SWITCH -> bottom, instant. Keyed on threadListItem.id + + * messages.length. `settledThreadIdRef` guards against streaming-token + * count ticks re-triggering the bottom-scroll. + * + * 2. NEW USER MESSAGE -> user pinned at top of viewport, smooth. + * Driven by the `isRunning` state transition (false -> true) via a + * post-commit effect rather than the `thread.runStart` event. + * WHY: assistant-ui emits `runStart` synchronously inside the runtime's + * repository mutation, BEFORE React commits the render that mounts the + * new rows. At that moment `virtualizerRef.current` can still be null + * (first message of a thread) and `viewportSize` can read 0. Waiting + * for the post-commit effect guarantees the Virtualizer is mounted and + * the new row is in the DOM. + * + * 3. STREAMING -> user owns scroll. Neither effect fires during a run. + */ const ThreadScrollController: FC = ({ virtualizerRef, captureBlankRef, @@ -177,6 +195,7 @@ const ThreadScrollController: FC = ({ const aui = useAui(); const threadId = useAuiState((s) => s.threadListItem.id); const messageCount = useAuiState((s) => s.thread.messages.length); + const isRunning = useAuiState((s) => s.thread.isRunning); const scrollToLastIndex = useCallback( (count: number, align: "start" | "end", smooth: boolean) => { @@ -188,8 +207,8 @@ const ThreadScrollController: FC = ({ [virtualizerRef], ); + // 1. Thread load/switch: snap to bottom once messages land. const settledThreadIdRef = useRef(null); - useEffect(() => { if (settledThreadIdRef.current !== threadId) { settledThreadIdRef.current = null; @@ -200,11 +219,23 @@ const ThreadScrollController: FC = ({ scrollToLastIndex(messageCount, "end", false); }, [threadId, messageCount, scrollToLastIndex]); - useAuiEvent("thread.runStart", () => { + // 2. isRunning false -> true: capture viewport blank + smooth-scroll the + // newly committed user message to the top. Effect-driven because by this + // point React has committed, the Virtualizer is mounted, and the row is + // in the DOM with a real measured height. `hasScrolledForRunRef` dedupes + // when message-count ticks during the run re-fire the effect. + const hasScrolledForRunRef = useRef(false); + useEffect(() => { + if (!isRunning) { + hasScrolledForRunRef.current = false; + return; + } + if (hasScrolledForRunRef.current) return; + hasScrolledForRunRef.current = true; captureBlankRef.current?.(); const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; scrollToLastIndex(liveCount, "start", true); - }); + }, [isRunning, aui, captureBlankRef, scrollToLastIndex]); return null; }; @@ -282,8 +313,10 @@ const VirtualizedMessages: FC = ({ const assistantMinHeight = streamingAssistantIndex >= 0 ? Math.max(0, blankSize - lastUserSize) : 0; - if (messageCount === 0) return null; - + // Always mount the Virtualizer — even when messageCount is 0 — so + // `virtualizerRef.current.viewportSize` is readable the moment the first + // message lands and `runStart` fires. The empty-state UI (ThreadWelcome / + // skeleton) is rendered by sibling AuiIf wrappers in Thread, not here. return ( Date: Sat, 18 Apr 2026 03:15:31 +0000 Subject: [PATCH 11/12] fix(thread): keep slack allocated after streaming ends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User reported the blank space below the assistant message collapsed as soon as streaming finished, bouncing the user message down from the top of the viewport. Root cause: both `blankSize` and the `isStreamingAssistant` check were gated on `isRunning`. When runStart -> runEnd transitioned, blankSize reset to 0 AND the index that owned the slack flipped to -1, so the minHeight style no longer applied. Fix: pin the slack to the specific message indexes that OWN the turn, not to the transient `isRunning` flag. - On captureBlank (fires from the isRunning false->true effect), snapshot the assistant index (count - 1) and user index (count - 2) into state. - `assistantMinHeight` keyed on `pinnedAssistantIndex >= 0`, so the slack persists after isRunning flips false. - Slack only releases when the pinned index no longer exists in the list (thread cleared, switched to shorter thread, or messages deleted), or when captureBlank runs again for the next turn. - `keepMounted` on the Virtualizer remains gated on `isRunning` — we only need to force-mount while the row is receiving streaming tokens. After streaming ends the row can unmount normally when scrolled away. Also read the count imperatively from `aui.thread().getState()` inside the capture (not from the useAuiState subscription) so the pinned indexes reflect the live post-commit count. --- src/components/assistant-ui/thread.tsx | 51 ++++++++++++++++++-------- 1 file changed, 35 insertions(+), 16 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 0a66198d..58d4a79a 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -265,31 +265,44 @@ const VirtualizedMessages: FC = ({ const aui = useAui(); const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); + // Slack state lives for the lifetime of an entire turn (not just while + // streaming). We pin it to the specific assistant index that owns the + // blank so the space stays allocated after isRunning flips false, and + // only gets released when the NEXT turn starts or the thread changes + // identity (count shrinks below the pinned index). const [blankSize, setBlankSize] = useState(0); const [lastUserSize, setLastUserSize] = useState(0); + const [pinnedAssistantIndex, setPinnedAssistantIndex] = useState(-1); + const [pinnedUserIndex, setPinnedUserIndex] = useState(-1); const [measuredUserElement, setMeasuredUserElement] = useState(null); useEffect(() => { captureBlankRef.current = () => { + // Read the current count imperatively; the subscription may not have + // ticked yet inside the isRunning effect that triggers this call. + const count = aui?.thread()?.getState()?.messages.length ?? 0; setBlankSize(virtualizerRef.current?.viewportSize ?? 0); + // Pin the indexes for this turn. The assistant row is the last row, + // the user row is the one before it. + setPinnedAssistantIndex(count - 1); + setPinnedUserIndex(count - 2); }; return () => { captureBlankRef.current = null; }; - }, [captureBlankRef, virtualizerRef]); + }, [aui, captureBlankRef, virtualizerRef]); + // Release slack when the pinned indexes no longer exist in the list + // (thread cleared, switched, or messages deleted). useEffect(() => { - if (!isRunning && blankSize !== 0) { + if (pinnedAssistantIndex >= messageCount && pinnedAssistantIndex !== -1) { setBlankSize(0); - } - }, [isRunning, blankSize]); - - useEffect(() => { - if (!isRunning && lastUserSize !== 0) { setLastUserSize(0); + setPinnedAssistantIndex(-1); + setPinnedUserIndex(-1); } - }, [isRunning, lastUserSize]); + }, [messageCount, pinnedAssistantIndex]); useEffect(() => { if (!measuredUserElement) return; @@ -308,10 +321,8 @@ const VirtualizedMessages: FC = ({ setMeasuredUserElement(el); }, []); - const streamingAssistantIndex = isRunning ? messageCount - 1 : -1; - const measuredUserIndex = isRunning ? messageCount - 2 : -1; const assistantMinHeight = - streamingAssistantIndex >= 0 ? Math.max(0, blankSize - lastUserSize) : 0; + pinnedAssistantIndex >= 0 ? Math.max(0, blankSize - lastUserSize) : 0; // Always mount the Virtualizer — even when messageCount is 0 — so // `virtualizerRef.current.viewportSize` is readable the moment the first @@ -321,24 +332,32 @@ const VirtualizedMessages: FC = ({ = 0 ? [streamingAssistantIndex] : undefined + isRunning && pinnedAssistantIndex >= 0 + ? [pinnedAssistantIndex] + : undefined } > {Array.from({ length: messageCount }, (_, index) => { const message = aui?.thread()?.getState()?.messages[index]; const id = message?.id ?? index; const role = message?.role; - const isStreamingAssistant = - role === "assistant" && index === streamingAssistantIndex; - const isMeasuredUser = role === "user" && index === measuredUserIndex; + const isPinnedAssistant = + role === "assistant" && index === pinnedAssistantIndex; + const isMeasuredUser = + role === "user" && index === pinnedUserIndex; return (
Date: Sat, 18 Apr 2026 03:25:31 +0000 Subject: [PATCH 12/12] fix(thread): correct scroll target + hover cleanup + thread-switch leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four issues, all valid from automated review: 1. Run-start scroll targeted the wrong row. At runStart the store holds [...prior, newUser, optimisticAssistant], so liveCount-1 is the assistant placeholder, not the user. We were scrolling the assistant row with align:'start', which placed the placeholder at the top and pushed the user message above the viewport — the exact opposite of the intended "question above answer" layout. Matches virtua's canonical chatbot demo: target `liveCount - 2` (the user) with align:'start', and let the assistant row's minHeight slack hold the user pinned at the top. 2. Thread-switch run-scroll race. Switching to an already-running thread would re-fire the top-anchor scroll, fighting the settled bottom scroll. Gate the run effect on `settledThreadIdRef.current === threadId` so it only fires for NEW runs started inside the currently-settled thread, not for thread switches. 3. Dedupe by threadId, not (threadId, messageCount). The previous key changed on every streaming token, re-triggering the scroll continuously during a stream. Refs should reset only when isRunning flips false. 4. Hover state leak on virtualization unmount. MessagePrimitive.Root (which we replaced) cleared isHovering in its ref cleanup; our hook only mirrored mouse events. When the virtualizer unmounted a hovered row without firing mouseleave, the message runtime stayed stuck in the hovered state and its ActionBar stayed visible when remounted. Capture the message runtime in a ref at hook time, track hovering state in a ref, and clear on unmount. 5. Slack leak across thread switches. If the new thread had >= pinned message count, the existing release effect (keyed on pinnedIndex >= messageCount) wouldn't fire, leaving stale minHeight on an unrelated row. Reset all slack state on threadId change explicitly. Also renamed scrollToLastIndex -> scrollToIndex and changed its signature to take an index directly (not count-1 derivation), because the caller now needs to target a non-last index. --- src/components/assistant-ui/thread.tsx | 99 +++++++++++++++++++------- 1 file changed, 75 insertions(+), 24 deletions(-) diff --git a/src/components/assistant-ui/thread.tsx b/src/components/assistant-ui/thread.tsx index 58d4a79a..239bcedb 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -197,11 +197,11 @@ const ThreadScrollController: FC = ({ const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); - const scrollToLastIndex = useCallback( - (count: number, align: "start" | "end", smooth: boolean) => { - if (count <= 0) return; + const scrollToIndex = useCallback( + (index: number, align: "start" | "end", smooth: boolean) => { + if (index < 0) return; requestAnimationFrame(() => { - virtualizerRef.current?.scrollToIndex(count - 1, { align, smooth }); + virtualizerRef.current?.scrollToIndex(index, { align, smooth }); }); }, [virtualizerRef], @@ -216,38 +216,78 @@ const ThreadScrollController: FC = ({ if (messageCount <= 0) return; if (settledThreadIdRef.current === threadId) return; settledThreadIdRef.current = threadId; - scrollToLastIndex(messageCount, "end", false); - }, [threadId, messageCount, scrollToLastIndex]); - - // 2. isRunning false -> true: capture viewport blank + smooth-scroll the - // newly committed user message to the top. Effect-driven because by this - // point React has committed, the Virtualizer is mounted, and the row is - // in the DOM with a real measured height. `hasScrolledForRunRef` dedupes - // when message-count ticks during the run re-fire the effect. - const hasScrolledForRunRef = useRef(false); + scrollToIndex(messageCount - 1, "end", false); + }, [threadId, messageCount, scrollToIndex]); + + // 2. isRunning false -> true for the ACTIVE thread: capture viewport blank + + // smooth-scroll the newly committed user message to the top. The user-row + // target matches virtua's chatbot demo exactly: scroll user to align:start, + // and the streaming assistant row beneath it carries minHeight slack so + // the user stays pinned even before assistant content streams in. + // + // Guard on `settledThreadIdRef.current === threadId` so that switching to + // an already-running thread doesn't re-fire this scroll (the settled + // effect above has already jumped to bottom). hasRunScrolledForThreadRef + // dedupes across streaming-token ticks within the same run. + const runScrolledForThreadRef = useRef(null); useEffect(() => { if (!isRunning) { - hasScrolledForRunRef.current = false; + runScrolledForThreadRef.current = null; return; } - if (hasScrolledForRunRef.current) return; - hasScrolledForRunRef.current = true; + // Only react once the thread-switch settling has completed for this id; + // prevents mount-time or thread-switch-to-running-thread re-scrolls. + if (settledThreadIdRef.current !== threadId) return; + // Dedupe per (thread, run): fire exactly once per false->true transition + // for this thread id. The ref resets when isRunning flips back to false, + // so the next run can fire again. + if (runScrolledForThreadRef.current === threadId) return; + runScrolledForThreadRef.current = threadId; captureBlankRef.current?.(); const liveCount = aui?.thread()?.getState()?.messages.length ?? 0; - scrollToLastIndex(liveCount, "start", true); - }, [isRunning, aui, captureBlankRef, scrollToLastIndex]); + // Target the USER row, not the last row. At runStart the store holds + // [...prior, newUser, optimisticAssistant], so the user is at liveCount-2. + // If only the user has landed (count=liveCount-1), fall back to count-1. + const userIndex = liveCount >= 2 ? liveCount - 2 : liveCount - 1; + scrollToIndex(userIndex, "start", true); + }, [ + isRunning, + threadId, + aui, + captureBlankRef, + scrollToIndex, + ]); return null; }; const useMessageHoverHandlers = () => { const aui = useAui(); + // Capture the message runtime at hook time so cleanup targets the same + // runtime even after the component unmounts (virtualization can remove a + // hovered row without ever firing mouseleave). Without this the message's + // isHovering state sticks true and its ActionBar stays pinned visible. + const hoveringRef = useRef(false); + const messageRef = useRef | null>(null); + messageRef.current = aui?.message() ?? null; + const onMouseEnter = useCallback(() => { - aui?.message()?.setIsHovering?.(true); - }, [aui]); + hoveringRef.current = true; + messageRef.current?.setIsHovering?.(true); + }, []); const onMouseLeave = useCallback(() => { - aui?.message()?.setIsHovering?.(false); - }, [aui]); + hoveringRef.current = false; + messageRef.current?.setIsHovering?.(false); + }, []); + + useEffect(() => { + return () => { + if (hoveringRef.current) { + messageRef.current?.setIsHovering?.(false); + } + }; + }, []); + return { onMouseEnter, onMouseLeave }; }; @@ -263,6 +303,7 @@ const VirtualizedMessages: FC = ({ captureBlankRef, }) => { const aui = useAui(); + const threadId = useAuiState((s) => s.threadListItem.id); const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); // Slack state lives for the lifetime of an entire turn (not just while @@ -293,8 +334,18 @@ const VirtualizedMessages: FC = ({ }; }, [aui, captureBlankRef, virtualizerRef]); - // Release slack when the pinned indexes no longer exist in the list - // (thread cleared, switched, or messages deleted). + // Release slack whenever the thread identity changes — otherwise stale + // minHeight can leak into a different thread if it happens to have + // enough messages to keep the pinned index valid. + useEffect(() => { + setBlankSize(0); + setLastUserSize(0); + setPinnedAssistantIndex(-1); + setPinnedUserIndex(-1); + }, [threadId]); + + // Also release when the pinned indexes no longer exist in the list + // (messages deleted, thread cleared). useEffect(() => { if (pinnedAssistantIndex >= messageCount && pinnedAssistantIndex !== -1) { setBlankSize(0);