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 e2efadc8..239bcedb 100644 --- a/src/components/assistant-ui/thread.tsx +++ b/src/components/assistant-ui/thread.tsx @@ -39,7 +39,7 @@ import { useMessagePartText, useAuiState, } from "@assistant-ui/react"; -import { useVirtualizer } from "@tanstack/react-virtual"; +import { Virtualizer, type VirtualizerHandle } from "virtua"; import type { FC, RefObject } from "react"; import { createContext, useContext } from "react"; @@ -115,8 +115,12 @@ interface ThreadProps { items?: Item[]; } +type ThreadVirtualizer = VirtualizerHandle; + export const Thread: FC = ({ items = [] }) => { const viewportRef = useRef(null); + const virtualizerRef = useRef(null); + const captureBlankRef = useRef<(() => void) | null>(null); return ( = ({ items = [] }) => { ["--thread-max-width" as string]: "50rem", }} > - thread.isLoading}> @@ -138,8 +146,16 @@ export const Thread: FC = ({ items = [] }) => { - - + + +
@@ -148,50 +164,263 @@ export const Thread: FC = ({ items = [] }) => { ); }; +interface ThreadScrollControllerProps { + virtualizerRef: RefObject; + 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, +}) => { + 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 scrollToIndex = useCallback( + (index: number, align: "start" | "end", smooth: boolean) => { + if (index < 0) return; + requestAnimationFrame(() => { + virtualizerRef.current?.scrollToIndex(index, { align, smooth }); + }); + }, + [virtualizerRef], + ); + + // 1. Thread load/switch: snap to bottom once messages land. + const settledThreadIdRef = useRef(null); + useEffect(() => { + if (settledThreadIdRef.current !== threadId) { + settledThreadIdRef.current = null; + } + if (messageCount <= 0) return; + if (settledThreadIdRef.current === threadId) return; + settledThreadIdRef.current = threadId; + 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) { + 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; + // 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; + // 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(() => { + hoveringRef.current = true; + messageRef.current?.setIsHovering?.(true); + }, []); + const onMouseLeave = useCallback(() => { + hoveringRef.current = false; + messageRef.current?.setIsHovering?.(false); + }, []); + + useEffect(() => { + return () => { + if (hoveringRef.current) { + messageRef.current?.setIsHovering?.(false); + } + }; + }, []); + + return { onMouseEnter, onMouseLeave }; +}; + interface VirtualizedMessagesProps { scrollRef: RefObject; + virtualizerRef: RefObject; + captureBlankRef: RefObject<(() => void) | null>; } -const VirtualizedMessages: FC = ({ scrollRef }) => { +const VirtualizedMessages: FC = ({ + scrollRef, + virtualizerRef, + 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 + // 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); - const virtualizer = useVirtualizer({ - count: messageCount, - getScrollElement: () => scrollRef.current, - estimateSize: () => 350, - overscan: 5, - }); + 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; + }; + }, [aui, captureBlankRef, virtualizerRef]); - if (messageCount === 0) return null; + // 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); + setLastUserSize(0); + setPinnedAssistantIndex(-1); + setPinnedUserIndex(-1); + } + }, [messageCount, pinnedAssistantIndex]); + + useEffect(() => { + if (!measuredUserElement) return; + const update = () => { + setLastUserSize(measuredUserElement.getBoundingClientRect().height); + }; + update(); + const resizeObserver = new ResizeObserver(update); + resizeObserver.observe(measuredUserElement); + return () => { + resizeObserver.disconnect(); + }; + }, [measuredUserElement]); + + const measureUserMessageRef = useCallback((el: HTMLDivElement | null) => { + setMeasuredUserElement(el); + }, []); + + const assistantMinHeight = + 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 + // message lands and `runStart` fires. The empty-state UI (ThreadWelcome / + // skeleton) is rendered by sibling AuiIf wrappers in Thread, not here. return ( -
= 0 + ? [pinnedAssistantIndex] + : undefined + } > - {virtualizer.getVirtualItems().map((virtualRow) => ( -
- -
- ))} -
+ {Array.from({ length: messageCount }, (_, index) => { + const message = aui?.thread()?.getState()?.messages[index]; + const id = message?.id ?? index; + const role = message?.role; + const isPinnedAssistant = + role === "assistant" && index === pinnedAssistantIndex; + const isMeasuredUser = + role === "user" && index === pinnedUserIndex; + + return ( +
+ +
+ ); + })} + ); }; @@ -1185,37 +1414,40 @@ const MessageError: FC = () => { }; const AssistantMessage: FC = () => { + const message = useMessage(); + const hover = useMessageHoverHandlers(); return ( - -
-
- - - -
+
+
+ + + +
-
- - -
+
+ +
- +
); }; @@ -1311,55 +1543,58 @@ const UserMessage: FC = () => { [expanded, showExpand], ); - return ( - -
- {/* Attachments display */} - - -
- - -
- - {showExpand && ( -
- -
- )} -
- -
+ const hover = useMessageHoverHandlers(); -
-
- + return ( +
+ {/* Attachments display */} + + +
+ + +
+ + {showExpand && ( +
+ +
+ )}
-
+ +
- +
+
+ +
- + + +
); };