Replace assistant-ui Viewport with custom scroll controller for virtualized thread - #394
Replace assistant-ui Viewport with custom scroll controller for virtualized thread#394urjitc wants to merge 12 commits into
Conversation
…e virtualizer conflicts
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 45 minutes and 43 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughReplaced TanStack Virtual with the Changes
Sequence DiagramsequenceDiagram
participant Client as Client/UI
participant TSC as ThreadScrollController
participant VRef as VirtualizerHandle
participant VComp as Virtualizer
participant Messages as Message Renderers
Client->>TSC: initialize / switch thread (threadId + messageCount)
TSC->>VRef: scrollToIndex(bottom)
VRef->>VComp: compute visible range
VComp->>Messages: mount/render visible rows
Client->>TSC: user message committed (run starts)
TSC->>TSC: detect viewport blank / defer during streaming
Messages->>TSC: streaming assistant updates (keepMounted for streaming index)
TSC->>VRef: scrollToIndex(userMessageIndex) (smooth)
VRef->>VComp: update range (rangeExtractor preserves streaming row)
VComp->>Messages: update DOM positions
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Greptile SummaryThis PR eliminates Two P2 concerns worth addressing before this pattern is extended:
Confidence Score: 4/5Safe to merge with awareness of two P2 scroll-accuracy issues that could cause subtle UX glitches. All findings are P2 (no data loss, no crashes, no auth issues). The score is 4 rather than 5 because one issue (stale handle count on runStart) is a latent correctness concern that could silently scroll to the wrong message and may be tricky to diagnose later, even though it doesn't break the primary path today. src/components/assistant-ui/thread.tsx — specifically the scrollIndexToTop fallback (lines 215–238) and the thread.runStart handler (lines 248–256). Important Files Changed
Reviews (1): Last reviewed commit: "Replace assistant-ui Viewport with custo..." | Re-trigger Greptile |
| (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], | ||
| ); |
There was a problem hiding this comment.
Manual offset estimation bypasses TanStack Virtual's own scroll API
getRowOffset only finds items in getVirtualItems() — the currently rendered window. When the target row is outside that window (common on first thread.runStart if the user scrolled up), found is undefined, and the code falls back to targetIndex * 350. This ignores any actual measured heights and can land the viewport on the wrong row.
TanStack Virtual exposes virtualizer.scrollToIndex(index, { align: 'start', behavior }) that handles unmeasured rows correctly using its internal size ledger. Exposing it through VirtualizerHandle removes the retry loop entirely:
// In VirtualizerHandle:
scrollToIndex: (index: number, options: { align?: 'start' | 'center' | 'end'; behavior?: ScrollBehavior }) => void;
// In VirtualizedMessages useEffect:
scrollToIndex: (index, options) => virtualizer.scrollToIndex(index, options),
// In ThreadScrollController – replace scrollIndexToTop with:
useAuiEvent("thread.runStart", () => {
const handle = virtualizerHandleRef.current;
if (!handle) return;
const targetIndex = handle.getCount() - 1;
handle.scrollToIndex(targetIndex, { align: 'start', behavior: 'smooth' });
});| 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"); | ||
| }); |
There was a problem hiding this comment.
handle.getCount() may reflect the pre-send count when thread.runStart fires
getCount is a closure over messageCount that is refreshed in a useEffect inside VirtualizedMessages. React runs sibling effects in tree order, but if useAuiEvent sets up its subscription via a layout effect (or the event fires synchronously during the action that mutates the AUI state), the VirtualizedMessages effect may not have committed yet. In that window getCount() returns the old value N, making targetIndex = N - 1 — the previous last message, not the new user message at index N.
A more robust approach is to read the count directly from AUI state at call time rather than from the stale closure:
useAuiEvent("thread.runStart", () => {
const handle = virtualizerHandleRef.current;
if (!handle) return;
// Read live count from AUI state to avoid stale handle closure
const liveCount = /* e.g., auiRef.current?.thread()?.getState()?.messages.length ?? */ handle.getCount();
const targetIndex = liveCount - 1;
scrollIndexToTop(targetIndex, "smooth");
});If you've verified that thread.runStart always fires after the handle update, adding a brief comment explaining that guarantee would prevent future confusion.
There was a problem hiding this comment.
2 issues found across 1 file
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/components/assistant-ui/thread.tsx">
<violation number="1" location="src/components/assistant-ui/thread.tsx:254">
P2: `thread.runStart` computes target index from potentially stale message count, which can intermittently scroll to the previous message instead of the newly sent one.</violation>
<violation number="2" location="src/components/assistant-ui/thread.tsx:302">
P2: `getRowOffset` searches `virtualizer.getVirtualItems()`, which only returns items in the visible viewport + overscan. When the target row is outside that window (e.g., user scrolled up before `thread.runStart`), the lookup returns `null` and the fallback `targetIndex * 350` ignores all previously measured row heights, landing the viewport on the wrong row.
TanStack Virtual already exposes `virtualizer.scrollToIndex(index, { align: 'start', behavior })` which uses its internal measurement cache and has built-in retry logic (up to 10 attempts) for dynamic sizing. Expose `scrollToIndex` through `VirtualizerHandle` instead of reimplementing scroll positioning manually.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
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.
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: Refactors core scroll and rendering logic for the chat thread, replacing library primitives with custom implementation. High UX impact risk requires human verification.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/components/assistant-ui/thread.tsx (1)
45-45: PreferRefObjectoverMutableRefObjectfor React 19 compatibility.Per React 19 updates,
useRef<T>(null)now returnsRefObject<T | null>.MutableRefObjectis deprecated in@types/react(though still exported in current 19.x versions). SincevirtualizerHandleRefis created viauseRef<VirtualizerHandle | null>(null)on line 121, its type is alreadyRefObject<VirtualizerHandle | null>. The explicitMutableRefObjectannotation on the props is redundant and inconsistent with the inferred type.♻️ Proposed change
-import type { FC, MutableRefObject, RefObject } from "react"; +import type { FC, RefObject } from "react";interface ThreadScrollControllerProps { scrollRef: RefObject<HTMLDivElement | null>; - virtualizerHandleRef: MutableRefObject<VirtualizerHandle | null>; + virtualizerHandleRef: RefObject<VirtualizerHandle | null>; }interface VirtualizedMessagesProps { scrollRef: RefObject<HTMLDivElement | null>; - virtualizerHandleRef: MutableRefObject<VirtualizerHandle | null>; + virtualizerHandleRef: RefObject<VirtualizerHandle | null>; }Also applies to: 185-188, 273-276
🤖 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 45, The import and prop typings use MutableRefObject but should use RefObject for React 19 compatibility: replace any use of MutableRefObject in the import and in the Thread component's prop types (e.g., the virtualizerHandleRef prop and other props currently typed as MutableRefObject) with RefObject (or remove redundant explicit typing where useRef already infers RefObject<VirtualizerHandle | null>), update the import to only include RefObject, and ensure the component signature and any usages (virtualizerHandleRef, and the other two occurrences noted) accept RefObject<T | null> to match useRef<VirtualizerHandle | null>(null).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/components/assistant-ui/thread.tsx`:
- Around line 204-253: The init/switch handlers call scrollToBottom which uses
el.scrollHeight (an estimated total) and can stop short; update the
thread.initialize and threadListItem.switchedTo handlers in
ThreadScrollController to instead read the live message count from
aui.thread().getState().messages.length, compute lastIndex = count - 1 (guarding
for count <= 0), and call virtualizerHandleRef.current?.scrollToIndex(lastIndex,
{ align: "end", behavior: "instant" }) (or "auto") inside requestAnimationFrame;
remove/stop using scrollToBottom for those two events so scrolling uses the
virtualizer's measured ledger like thread.runStart does.
---
Nitpick comments:
In `@src/components/assistant-ui/thread.tsx`:
- Line 45: The import and prop typings use MutableRefObject but should use
RefObject for React 19 compatibility: replace any use of MutableRefObject in the
import and in the Thread component's prop types (e.g., the virtualizerHandleRef
prop and other props currently typed as MutableRefObject) with RefObject (or
remove redundant explicit typing where useRef already infers
RefObject<VirtualizerHandle | null>), update the import to only include
RefObject, and ensure the component signature and any usages
(virtualizerHandleRef, and the other two occurrences noted) accept RefObject<T |
null> to match useRef<VirtualizerHandle | null>(null).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d3258e7f-d853-40e4-a54f-5ee4054c4c82
📒 Files selected for processing (1)
src/components/assistant-ui/thread.tsx
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).
- 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<typeof defaultRangeExtractor>[0] with the exported Range type. - Collapse verbose JSDoc and inline comments that duplicated adjacent code. Net -35 lines.
There was a problem hiding this comment.
3 issues found across 1 file (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/components/assistant-ui/thread.tsx">
<violation number="1" location="src/components/assistant-ui/thread.tsx:294">
P1: `.map()` always creates a new array reference, so `useAuiState` (which uses `Object.is` comparison) sees a "new" value on every store update and re-renders the component — including on every streaming token. This defeats the stated goal of the two-subscription split.
Consider subscribing to `s.thread.messages` directly (which is likely a stable reference that only changes when messages are added/removed) and deriving IDs in the component body, or splitting into `useAuiState((s) => s.thread.messages.length)` for the count and a separate stable mechanism for keys.</violation>
<violation number="2" location="src/components/assistant-ui/thread.tsx:327">
P2: Custom agent: **Flag AI Slop and Fabricated Changes**
Comment references `ROW_ESTIMATE_SIZE` (350) but the code assigns `SCROLL_PADDING_START` (16) — the wrong constant name was hallucinated. The same explanation also already exists at the `SCROLL_PADDING_START` constant definition, making this a duplicate narrating comment with an incorrect reference. Remove or correct the inline comment.</violation>
<violation number="3" location="src/components/assistant-ui/thread.tsx:367">
P1: Wrapper-only translate breaks positioning when `rangeExtractor` injects non-contiguous rows (pinned streaming row).</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
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.
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: This is a significant refactor of the core thread scrolling and virtualization logic. Changes to scroll behavior and layout observers in a main UI component require human review.
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.
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: This is a significant refactor of the thread's scrolling and virtualization logic. It replaces library primitives with custom logic, which requires human validation for UX consistency.
…tion 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.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/components/assistant-ui/thread.tsx (1)
347-351: Apply padding slack conditionally during active runs.
paddingEnd: 800pxserves a real purpose while anchoring a just-sent user message at the viewport top, but unconditionally applying it leaves a large empty scrollable region in idle/history views. The slack should only apply when streaming is active.♻️ Proposed adjustment
- paddingEnd: BOTTOM_SLACK_PX, + paddingEnd: isRunning ? BOTTOM_SLACK_PX : 0,The
isRunningvariable is already available in scope (line 312) and used elsewhere in the component for streaming-specific behavior.🤖 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 347 - 351, The unconditional bottom padding (paddingEnd: BOTTOM_SLACK_PX) is causing a large empty scroll region; change it to only apply when streaming by using the existing isRunning flag (e.g., set paddingEnd to BOTTOM_SLACK_PX when isRunning is true, otherwise 0). Locate the object/property that sets paddingEnd (refer to paddingEnd and BOTTOM_SLACK_PX in this component) and replace the constant with a conditional expression using isRunning so idle/history views have no extra slack while active runs retain the anchor padding.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/components/assistant-ui/thread.tsx`:
- Around line 289-297: useMessageHoverHandlers currently sets hover on mouse
enter/leave but doesn't clear state on component unmount which can leave
ActionBar stuck; add an effect cleanup inside useMessageHoverHandlers that calls
aui?.message()?.setIsHovering?.(false) on unmount (use useEffect with aui as
dependency) so any hovered state is cleared if the row unmounts, while retaining
the existing onMouseEnter/onMouseLeave handlers and guarding the call the same
way as in those callbacks.
- Around line 316-344: Subscribe VirtualizedMessages to the current thread
identity (e.g. add const threadId = useAuiState(s => s.thread.id) or
threadListItem.id) and use that value to namespace the virtualizer keys and
dependencies so measurements reset on thread switches: include threadId in the
rangeExtractor dependency array and change getItemKey to return a namespaced key
like `${threadId}:${aui?.thread()?.getState()?.messages[index]?.id ?? index}`
(so the virtualizer sees a new logical dataset when threadId changes).
---
Nitpick comments:
In `@src/components/assistant-ui/thread.tsx`:
- Around line 347-351: The unconditional bottom padding (paddingEnd:
BOTTOM_SLACK_PX) is causing a large empty scroll region; change it to only apply
when streaming by using the existing isRunning flag (e.g., set paddingEnd to
BOTTOM_SLACK_PX when isRunning is true, otherwise 0). Locate the object/property
that sets paddingEnd (refer to paddingEnd and BOTTOM_SLACK_PX in this component)
and replace the constant with a conditional expression using isRunning so
idle/history views have no extra slack while active runs retain the anchor
padding.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9b63f040-3467-4cec-9ec1-c4b09504096a
📒 Files selected for processing (1)
src/components/assistant-ui/thread.tsx
| 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 — 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], | ||
| ); | ||
|
|
||
| const virtualizer = useVirtualizer({ | ||
| const virtualizer = useVirtualizer<HTMLDivElement, HTMLDivElement>({ | ||
| count: messageCount, | ||
| getScrollElement: () => scrollRef.current, | ||
| estimateSize: () => 350, | ||
| overscan: 5, | ||
| 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, |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Confirm whether VirtualizedMessages observes thread identity and whether keys are thread-scoped.
sed -n '306,365p' src/components/assistant-ui/thread.tsx
printf '\n--- thread id / getItemKey / measure reset references ---\n'
rg -n -C3 'threadListItem\.id|getItemKey|virtualizer\.measure\(' src/components/assistant-ui/thread.tsxRepository: ThinkEx-OSS/thinkex
Length of output: 4680
Subscribe to thread identity in virtualizer to reset measurements on thread switches.
VirtualizedMessages only subscribes to messages.length and isRunning. When switching threads with the same message count and running state, React won't re-render, leaving TanStack Virtual's measurement cache and item keys from the previous thread intact. This breaks measurement accuracy for the new thread's messages.
Since ThreadScrollController already subscribes to threadListItem.id to handle thread switches, add the same subscription to VirtualizedMessages and namespace getItemKey by thread identity so the virtualizer recognizes a new logical dataset:
const messageCount = useAuiState((s) => s.thread.messages.length);
const isRunning = useAuiState((s) => s.thread.isRunning);
+ const threadId = useAuiState((s) => s.threadListItem.id);
// ...
getItemKey: (index) =>
- aui?.thread()?.getState()?.messages[index]?.id ?? index,
+ `${threadId ?? "thread"}:${
+ aui?.thread()?.getState()?.messages[index]?.id ?? index
+ }`,🤖 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 316 - 344, Subscribe
VirtualizedMessages to the current thread identity (e.g. add const threadId =
useAuiState(s => s.thread.id) or threadListItem.id) and use that value to
namespace the virtualizer keys and dependencies so measurements reset on thread
switches: include threadId in the rangeExtractor dependency array and change
getItemKey to return a namespaced key like
`${threadId}:${aui?.thread()?.getState()?.messages[index]?.id ?? index}` (so the
virtualizer sees a new logical dataset when threadId changes).
There was a problem hiding this comment.
1 issue found across 1 file (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/components/assistant-ui/thread.tsx">
<violation number="1" location="src/components/assistant-ui/thread.tsx:267">
P2: First message in an empty/switched thread can trigger both bottom-scroll and top-anchor scroll paths, causing conflicting auto-scroll behavior.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
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.
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
…art event 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.
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.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
src/components/assistant-ui/thread.tsx (2)
243-251:⚠️ Potential issue | 🟡 MinorClear hover state when virtualized rows unmount.
A hovered row can unmount without firing
onMouseLeave, leavingmessage.setIsHovering(true)stuck for that message.🧹 Proposed fix
const useMessageHoverHandlers = () => { const aui = useAui(); const onMouseEnter = useCallback(() => { aui?.message()?.setIsHovering?.(true); }, [aui]); const onMouseLeave = useCallback(() => { aui?.message()?.setIsHovering?.(false); }, [aui]); + useEffect(() => { + return () => { + aui?.message()?.setIsHovering?.(false); + }; + }, [aui]); return { onMouseEnter, onMouseLeave }; };🤖 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 243 - 251, The hover state can remain true if a virtualized row unmounts without onMouseLeave; update useMessageHoverHandlers to include a useEffect cleanup that calls aui?.message()?.setIsHovering?.(false) on unmount so the hover flag is cleared, while keeping the existing onMouseEnter/onMouseLeave callbacks (referencing useMessageHoverHandlers and aui.message().setIsHovering).
266-352:⚠️ Potential issue | 🟠 MajorSubscribe to thread identity so same-count thread switches re-render.
VirtualizedMessagesonly subscribes tomessageCountandisRunning. Switching to another thread with the same message count/running state can leave the old rows rendered because Line 329 reads messages imperatively only during render. Also key the virtualizer/rows by thread to resetvirtuameasurements for the new dataset.🐛 Proposed fix
const aui = useAui(); const messageCount = useAuiState((s) => s.thread.messages.length); const isRunning = useAuiState((s) => s.thread.isRunning); + const threadId = useAuiState((s) => s.threadListItem.id); @@ <Virtualizer + key={threadId ?? "thread"} ref={virtualizerRef} scrollRef={scrollRef} keepMounted={ @@ return ( <div - key={id} + key={`${threadId ?? "thread"}:${id}`} ref={isMeasuredUser ? measureUserMessageRef : undefined}🤖 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 266 - 352, The component only subscribes to messageCount and isRunning so switching to a different thread with the same counts leaves old rows; fix by subscribing to the thread identity and using it to drive re-renders and reset virtualizer measurements: read the thread id and messages from state (e.g. add useAuiState selectors that return s.thread.id and s.thread.messages or a combined {id,messages}) instead of imperatively calling aui.thread().getState() inside the render, and add the thread id as the key on the Virtualizer (and/or the outer row wrapper) so Virtualizer and rows remount when the thread changes (update references to Virtualizer, messageCount, streamingAssistantIndex, measuredUserIndex, and the row mapping to use the selected messages/thread id).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@package.json`:
- Line 96: package.json now includes the new dependency "virtua" but the
generated pnpm lockfile is missing; run pnpm install locally to produce an
updated pnpm-lock.yaml, add and commit that pnpm-lock.yaml alongside the
package.json change so installs are reproducible and CI uses the same locked
dependency versions.
In `@src/components/assistant-ui/thread.tsx`:
- Around line 200-238: The effect that runs when isRunning flips currently calls
scrollToLastIndex(liveCount, "start", true) which targets the final row
(assistant placeholder); instead compute the user-row count and scroll to that
instead — e.g. calculate liveCount = aui?.thread()?.getState()?.messages.length
?? 0, if liveCount <= 1 return, then call scrollToLastIndex(liveCount - 1,
"start", true) (since scrollToLastIndex(count) scrolls to index count-1) so the
user message (messageCount-2 index) is scrolled into view; update the isRunning
effect that references scrollToLastIndex, hasScrolledForRunRef, and
captureBlankRef accordingly.
---
Duplicate comments:
In `@src/components/assistant-ui/thread.tsx`:
- Around line 243-251: The hover state can remain true if a virtualized row
unmounts without onMouseLeave; update useMessageHoverHandlers to include a
useEffect cleanup that calls aui?.message()?.setIsHovering?.(false) on unmount
so the hover flag is cleared, while keeping the existing
onMouseEnter/onMouseLeave callbacks (referencing useMessageHoverHandlers and
aui.message().setIsHovering).
- Around line 266-352: The component only subscribes to messageCount and
isRunning so switching to a different thread with the same counts leaves old
rows; fix by subscribing to the thread identity and using it to drive re-renders
and reset virtualizer measurements: read the thread id and messages from state
(e.g. add useAuiState selectors that return s.thread.id and s.thread.messages or
a combined {id,messages}) instead of imperatively calling
aui.thread().getState() inside the render, and add the thread id as the key on
the Virtualizer (and/or the outer row wrapper) so Virtualizer and rows remount
when the thread changes (update references to Virtualizer, messageCount,
streamingAssistantIndex, measuredUserIndex, and the row mapping to use the
selected messages/thread id).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a6c95687-2bcf-44c4-b8e5-0a8624801a28
📒 Files selected for processing (2)
package.jsonsrc/components/assistant-ui/thread.tsx
| "@tanstack/react-query": "^5.96.1", | ||
| "@tanstack/react-query-devtools": "^5.96.1", | ||
| "@tanstack/react-virtual": "^3.13.23", | ||
| "virtua": "^0.49.1", |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Description: Verify pnpm lockfile reflects the package.json virtualizer change.
# Expected: pnpm-lock.yaml exists, contains virtua, and does not contain `@tanstack/react-virtual`.
test -f pnpm-lock.yaml
rg -n '"?virtua"?|/virtua@' pnpm-lock.yaml
if rg -n '@tanstack/react-virtual' pnpm-lock.yaml; then
printf 'Stale `@tanstack/react-virtual` lockfile entry found\n' >&2
exit 1
fiRepository: ThinkEx-OSS/thinkex
Length of output: 45
🏁 Script executed:
#!/bin/bash
# Check if pnpm-lock.yaml exists and examine virtualizer entries
echo "=== Checking pnpm-lock.yaml existence ==="
if [ -f pnpm-lock.yaml ]; then
echo "✓ pnpm-lock.yaml exists"
echo ""
echo "=== Searching for 'virtua' entries ==="
rg -i 'virtua' pnpm-lock.yaml | head -20
echo ""
echo "=== Searching for '@tanstack/react-virtual' entries ==="
rg '@tanstack/react-virtual' pnpm-lock.yaml | head -20
else
echo "✗ pnpm-lock.yaml does NOT exist"
fi
echo ""
echo "=== Checking package.json for virtua ==="
rg 'virtua' package.jsonRepository: ThinkEx-OSS/thinkex
Length of output: 203
Add pnpm-lock.yaml to the commit.
The virtua dependency has been added to package.json but pnpm-lock.yaml is missing from the repository. Projects using pnpm require the lockfile to ensure reproducible installs. Run pnpm install locally to generate and commit the updated lockfile.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@package.json` at line 96, package.json now includes the new dependency
"virtua" but the generated pnpm lockfile is missing; run pnpm install locally to
produce an updated pnpm-lock.yaml, add and commit that pnpm-lock.yaml alongside
the package.json change so installs are reproducible and CI uses the same locked
dependency versions.
There was a problem hiding this comment.
2 issues found across 1 file (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/components/assistant-ui/thread.tsx">
<violation number="1" location="src/components/assistant-ui/thread.tsx:228">
P2: Run-scroll effect is not guarded to same-thread `false -> true` transition, so it can trigger on mount/thread switch for already-running threads and cause unintended scroll jumps.</violation>
<violation number="2" location="src/components/assistant-ui/thread.tsx:299">
P2: Pinned slack state is not reset on thread identity changes, so stale min-height can leak into another thread when message count remains high enough.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
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.
|
You're iterating quickly on this pull request. To help protect your rate limits, cubic has paused automatic reviews on new pushes for now—when you're ready for another review, comment |
| // effect above has already jumped to bottom). hasRunScrolledForThreadRef | ||
| // dedupes across streaming-token ticks within the same run. | ||
| const runScrolledForThreadRef = useRef<string | null>(null); | ||
| useEffect(() => { |
There was a problem hiding this comment.
[🟡 Medium] [🔵 Bug]
The second scroll effect is keyed on isRunning being truthy, not on a false -> true transition. On a thread switch into a thread that is already running, React runs the earlier settling effect first, which sets settledThreadIdRef.current = threadId; this effect then runs in the same commit and calls scrollToIndex(userIndex, "start", true), overriding the intended bottom snap for thread switches. The result is that opening an in-progress thread jumps back to the current turn's user message instead of showing the newest assistant output. Gate this effect on an actual run-start transition for the active thread (for example by tracking the previous isRunning/threadId pair) rather than isRunning truthiness alone.
// src/components/assistant-ui/thread.tsx
useEffect(() => {
if (!isRunning) {
runScrolledForThreadRef.current = null;
return;
}| } | ||
| // 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; |
There was a problem hiding this comment.
[🟡 Medium] [🔵 Bug]
ThreadScrollController intends to suppress the smooth align:"start" scroll when the user switches into a thread that is already mid-run, but the two effects execute in the opposite way: the load/switch effect marks the new threadId as settled first, then the isRunning effect immediately passes its guard and scrolls to the last user row anyway. That means opening a thread with an active assistant response will jump the viewport away from the bottom and pin the previous user message at the top, which is the exact scenario the comment says should not happen. The fix is to distinguish a real local false -> true run transition from a plain thread switch (for example by tracking previous isRunning/threadId together, or by skipping the run-scroll on the first committed render after a thread change).
// src/components/assistant-ui/thread.tsx
if (settledThreadIdRef.current !== threadId) return;
if (runScrolledForThreadRef.current === threadId) return;
runScrolledForThreadRef.current = threadId;
captureBlankRef.current?.();
const liveCount = 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); |
There was a problem hiding this comment.
[🟡 Medium] [🔵 Bug]
The controller explicitly handles a run-start state where only the new user row has committed, but captureBlankRef still hard-codes count - 1 as the assistant row and count - 2 as the user row. In that count === 1 case it pins index 0 as the assistant and -1 as the user, and there is no follow-up effect that repairs those indices when the optimistic assistant row is appended. As a result, keepMounted/minHeight attach to the wrong item and the first turn (or any delayed assistant append) loses the slack that is supposed to keep the user message anchored at the top while streaming.
// src/components/assistant-ui/thread.tsx
const count = aui?.thread()?.getState()?.messages.length ?? 0;
setBlankSize(virtualizerRef.current?.viewportSize ?? 0);
setPinnedAssistantIndex(count - 1);
setPinnedUserIndex(count - 2);Derive the pinned indices from the current message roles, or detect the one-row case and move the assistant pin forward when the placeholder row appears.
| // 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; |
There was a problem hiding this comment.
[🟡 Medium] [🔵 Bug]
The new virtua scroll controller assumes every isRunning transition appends a [newUser, optimisticAssistant] pair and derives its target from positional math, but assistant-ui’s branching behavior also starts runs when a user message is edited or an assistant message is reloaded. In those cases the message list is branched in place rather than gaining a fresh user row, so this code scrolls to the penultimate existing message and pairs the blank-space pinning with the wrong rows, producing an unexpected viewport jump during reload/edit runs. Gate the smooth-scroll logic on an actual appended user message (or derive the target from the message that started the run) instead of hard-coding liveCount - 2 / count - 1.
// src/components/assistant-ui/thread.tsx
const liveCount = aui?.thread()?.getState()?.messages.length ?? 0;
// Target the USER row, not the last row. At runStart the store holds
// [...prior, newUser, optimisticAssistant], so the user is at liveCount-2.
const userIndex = liveCount >= 2 ? liveCount - 2 : liveCount - 1;
scrollToIndex(userIndex, "start", true);
Removes
ThreadPrimitive.ViewportandMessagePrimitive.Rootfrom the virtualized thread to eliminate ResizeObserver/MutationObserver thrashing and Slack min-height injection that fight@tanstack/react-virtual. Owns scroll behavior explicitly with a lightweightThreadScrollControllerthat hooks intothread.initialize,threadListItem.switchedTo, andthread.runStartevents.ThreadPrimitive.Viewportwith plain<div>scroll container; adduseAuiEventfrom@assistant-ui/reactfor scroll event handling; implementThreadScrollControllerthat scrolls to bottom on thread init/switch (instant) and anchors the new user message at the top onthread.runStart(smooth); removeMessagePrimitive.Rootwrappers, replace with direct divs anduseMessageHoverHandlersto preserve ActionBar autohide viamessage.setIsHovering; exposeVirtualizerHandlefromVirtualizedMessagesso scroll controller can translate indices into virtual row offsets; preserve all existing styling, loading states, and composer logic.Summary by CodeRabbit