wip: scroll anchoring improvements for chat

packages/@react-spectrum/ai/src/ResponseStatus.tsx

@@ -670,7 +670,10 @@ const executionTraceItemStyles = style({
       default: 'block',
       ':last-child': 'none'
     }
-  },
+  }
+});
+
+const executionTraceItemEntranceStyles = style({
   transition: '[opacity, translate]',
   transitionDuration: `[${EXECUTION_TRACE_ITEM_TRANSITION_DURATION}ms, 310ms]`,
   transitionTimingFunction: `[cubic-bezier(0.45, 0, 0.4, 1), ${EXECUTION_TRACE_ITEM_TIMING_FUNCTION}]`,
@@ -735,9 +738,27 @@ export const ExecutionTraceItem = forwardRef(function ExecutionTraceItem(
   let domProps = filterDOMProps(otherProps);
   let {isFocusVisible, focusProps} = useFocusRing();
   let hasDetail = detail != null;
+  // Play the entrance (fade + slide) once, then remove the animating styles.
+  // This is to prevent a flash that occurs when scrolling in virtualized containers
+  // because the browser keeps re-creating the layer these animations force it onto
+  let [hasEntered, setHasEntered] = useState(false);
 
   return (
-    <li {...domProps} ref={domRef} className={mergeStyles(executionTraceItemStyles, styles)}>
+    <li
+      {...domProps}
+      ref={domRef}
+      onTransitionEnd={e => {
+        // Only react to this item's own opacity transition (the longer of the two, so both the
+        // fade and slide have finished), not transitions bubbling up from descendants.
+        if (e.target === e.currentTarget && e.propertyName === 'opacity') {
+          setHasEntered(true);
+        }
+      }}
+      className={mergeStyles(
+        executionTraceItemStyles,
+        hasEntered ? undefined : executionTraceItemEntranceStyles,
+        styles
+      )}>
       <div className={executionTraceItemIconContainerStyles}>
         <CenterBaseline>
           {status === 'failed' && (

packages/@react-spectrum/ai/stories/Chat.stories.tsx

@@ -62,7 +62,7 @@ const meta: Meta<typeof Chat> = {
   title: 'AI/Chat',
   decorators: [
     Story => (
-      <div style={{width: '800px', height: '600px'}}>
+      <div style={{width: '80vw', height: '600px'}}>
         <Story />
       </div>
     )

packages/react-stately/src/virtualizer/ScrollAnchor.ts

@@ -88,22 +88,48 @@ export function captureScrollAnchor(
   isAnchorable: (layoutInfo: LayoutInfo) => boolean = () => true
 ): ScrollAnchor | null {
   let dimension = dimensionForAxis(axis);
+  // The corner on the item's leading edge - the side where content is added/removed. For 'end'
+  // that's the start of the axis (top/left); for 'start' it's the end of the axis (bottom/right).
+  let corner: RectCorner =
+    axis === 'x'
+      ? edge === 'end'
+        ? 'topLeft'
+        : 'topRight'
+      : edge === 'end'
+        ? 'topLeft'
+        : 'bottomLeft';
+  let viewportExtent = visibleRect[dimension];
   let best: ScrollAnchor | null = null;
+  // Fallback used only when every visible item is clipped past the leading edge (e.g. a single item
+  // taller than the viewport): the least-clipped item still makes the most stable anchor available.
+  let fallback: ScrollAnchor | null = null;
   for (let [key, layoutInfo] of visibleLayoutInfos) {
     if (!layoutInfo || !isAnchorable(layoutInfo)) {
       continue;
     }
     let overlap = layoutInfo.rect.intersection(visibleRect)[dimension];
-    if (layoutInfo.rect.area > 0 && overlap >= MIN_ANCHOR_OVERLAP) {
-      let corner = layoutInfo.rect.getCornerInRect(visibleRect) ?? 'topLeft';
-      let offset = layoutInfo.rect[corner][axis] - visibleRect[axis];
+    if (layoutInfo.rect.area <= 0 || overlap < MIN_ANCHOR_OVERLAP) {
+      continue;
+    }
+    let offset = layoutInfo.rect[corner][axis] - visibleRect[axis];
+    // Is the leading edge within the viewport?
+    let leadingEdgeVisible = edge === 'end' ? offset >= 0 : offset <= viewportExtent;
+    if (leadingEdgeVisible) {
+      // Pick the item nearest the leading edge among those whose leading edge is visible.
       let isBetter = !best || (edge === 'end' ? offset < best.offset : offset > best.offset);
       if (isBetter) {
         best = {key, corner, offset};
       }
+    } else {
+      // Least-clipped = closest to the leading edge from the clipped side.
+      let isBetter =
+        !fallback || (edge === 'end' ? offset > fallback.offset : offset < fallback.offset);
+      if (isBetter) {
+        fallback = {key, corner, offset};
+      }
     }
   }
-  return best;
+  return best ?? fallback;
 }
 
 /** Returns the viewport coordinate (along `axis`) that pins the viewport to `edge` of the content. */
@@ -154,7 +180,8 @@ export function resolveScrollAdjustment(
   contentSizeDelta: number,
   getLayoutInfo: (key: Key) => LayoutInfo | null,
   previousVisibleRect: Rect,
-  contentSize: Size
+  contentSize: Size,
+  changeIsAtEdge: boolean = true
 ): Rect | null {
   let withTarget = (target: number): Rect =>
     axis === 'x'
@@ -171,7 +198,16 @@ export function resolveScrollAdjustment(
           previousVisibleRect.height
         );
 
-  if (anchor) {
+  // Two possible responses when content settles: "preserve the anchor" (keep the item the user is
+  // looking at in place) or "follow the edge" (keep the viewport pinned to the
+  // content edge, e.g. the bottom of a chat)
+  let followEdge =
+    wasNearAnchorEdge &&
+    !isScrolling &&
+    itemSizeChanged &&
+    contentSizeDelta !== 0 &&
+    changeIsAtEdge;
+  if (anchor && !followEdge) {
     let target = computeScrollAnchorTarget(
       anchor,
       axis,
@@ -184,7 +220,12 @@ export function resolveScrollAdjustment(
     }
   }
 
-  if (wasNearAnchorEdge && !isScrolling && (!itemSizeChanged || contentSizeDelta > 0)) {
+  if (
+    wasNearAnchorEdge &&
+    !isScrolling &&
+    (!itemSizeChanged || contentSizeDelta !== 0) &&
+    changeIsAtEdge
+  ) {
     let target = withTarget(getEdgeSnapTarget(edge, axis, contentSize, previousVisibleRect));
     return target.equals(previousVisibleRect) ? null : target;
   }
@@ -196,14 +237,18 @@ export interface ResolveAfterLayoutOptions {
   anchorInfo: ScrollAnchorInfo | null;
   /** The anchor captured by `captureBeforeLayout` before this pass's `layout.update()` ran. */
   anchor: ScrollAnchor | null;
-  /** The full post-layout visible layout infos, i.e. `virtualizer.getVisibleLayoutInfos()`. */
-  postLayoutInfos: Map<Key, LayoutInfo>;
   previousVisibleRect: Rect;
   previousContentSize: Size;
   contentSize: Size;
   itemSizeChanged: boolean;
   isScrolling: boolean;
   getLayoutInfo: (key: Key) => LayoutInfo | null;
+  /**
+   * Whether the content that changed this pass was at the anchored edge (e.g. the newest item in a
+   * bottom-anchored list). When false, the viewport does not follow the edge, so a mid-list resize
+   * while the user is scrolled away preserves their position. Defaults to true.
+   */
+  changeIsAtEdge?: boolean;
 }
 
 /**
@@ -212,14 +257,10 @@ export interface ResolveAfterLayoutOptions {
  */
 export class ScrollAnchorTracker {
   private hasSnappedToEdge = false;
-  private hadEstimatedVisibleItems = false;
-  private wasNearAnchorEdge = false;
 
-  /** Resets all tracked state, e.g. when the virtualizer's layout instance changes. */
+  /** Resets the first-layout flag, e.g. when the virtualizer's layout instance changes. */
   reset(): void {
     this.hasSnappedToEdge = false;
-    this.hadEstimatedVisibleItems = false;
-    this.wasNearAnchorEdge = false;
   }
 
   /**
@@ -250,43 +291,19 @@ export class ScrollAnchorTracker {
     let {
       anchorInfo,
       anchor,
-      postLayoutInfos,
       previousVisibleRect,
       previousContentSize,
       contentSize,
       itemSizeChanged,
       isScrolling,
-      getLayoutInfo
+      getLayoutInfo,
+      changeIsAtEdge = true
     } = options;
 
     if (!anchorInfo) {
       return null;
     }
 
-    // Read the previous pass's state into locals before any writes below overwrite it.
-    let wasSettlingLastPass = this.hadEstimatedVisibleItems;
-    let wasNearAnchorEdgeLastPass = this.wasNearAnchorEdge;
-
-    let hasEstimated = false;
-    for (let layoutInfo of postLayoutInfos.values()) {
-      if (layoutInfo.estimatedSize) {
-        hasEstimated = true;
-        break;
-      }
-    }
-    this.hadEstimatedVisibleItems = hasEstimated;
-    // Don't recheck "near edge?" mid-resize because it could look like a scroll that never happened.
-    // Reuse the answer from before the resizing started.
-    if (!wasSettlingLastPass) {
-      this.wasNearAnchorEdge = isNearEdge(
-        previousVisibleRect,
-        previousContentSize,
-        anchorInfo.edge,
-        anchorInfo.axis,
-        anchorInfo.threshold
-      );
-    }
-
     if (previousVisibleRect.area === 0) {
       return null;
     }
@@ -303,23 +320,16 @@ export class ScrollAnchorTracker {
 
     let wasNearAnchorEdge =
       isFirstAnchoredLayout ||
-      (wasSettlingLastPass && wasNearAnchorEdgeLastPass) ||
       isNearEdge(
         previousVisibleRect,
         previousContentSize,
         anchorInfo.edge,
         anchorInfo.axis,
         anchorInfo.threshold
       );
-    // A first-ever layout always snaps to the edge, even if the raw distance check says
-    // otherwise. Save that real decision here so later passes in this cascade reuse it.
-    if (!wasSettlingLastPass) {
-      this.wasNearAnchorEdge = wasNearAnchorEdge;
-    }
-    // Skip restoring to the captured anchor while still resizing because items above it are also still growing,
-    // and following it would fall short of the edge.
-    let effectiveAnchor =
-      isFirstAnchoredLayout || (wasSettlingLastPass && wasNearAnchorEdgeLastPass) ? null : anchor;
+    let effectiveAnchor = isFirstAnchoredLayout ? null : anchor;
+    // The first anchored layout always snaps to the edge, regardless of what changed.
+    let effectiveChangeIsAtEdge = isFirstAnchoredLayout || changeIsAtEdge;
     return resolveScrollAdjustment(
       anchorInfo.edge,
       anchorInfo.axis,
@@ -330,7 +340,8 @@ export class ScrollAnchorTracker {
       contentSizeDelta,
       getLayoutInfo,
       previousVisibleRect,
-      contentSize
+      contentSize,
+      effectiveChangeIsAtEdge
     );
   }
 }

packages/react-stately/src/virtualizer/Virtualizer.ts

@@ -71,7 +71,12 @@ export class Virtualizer<T extends object, V> {
   private _isScrolling: boolean;
   private _invalidationContext: InvalidationContext;
   private _overscanManager: OverscanManager;
+
   private _scrollAnchor: ScrollAnchorTracker;
+  // Together they classify whether the changed content was at the anchored (e.g. bottom) edge to avoid
+  // following the edge when only a mid-list item resized in a short chat.
+  private _hadItemResize: boolean;
+  private _batchIncludedNewestContent: boolean;
 
   constructor(options: VirtualizerOptions<T, V>) {
     this.delegate = options.delegate;
@@ -88,6 +93,8 @@ export class Virtualizer<T extends object, V> {
     this._invalidationContext = {};
     this._overscanManager = new OverscanManager();
     this._scrollAnchor = new ScrollAnchorTracker();
+    this._hadItemResize = false;
+    this._batchIncludedNewestContent = false;
   }
 
   /** Returns whether the given key, or an ancestor, is persisted. */
@@ -177,14 +184,11 @@ export class Virtualizer<T extends object, V> {
     // On first render _visibleViews is empty so no anchor will be found.
     let anchor: ScrollAnchor | null = null;
     if (anchorInfo) {
-      let preLayoutInfos: [Key, LayoutInfo][] = [];
-      for (let [key, view] of this._visibleViews) {
-        let layoutInfo = this.layout.getLayoutInfo(key) ?? view.layoutInfo;
-        if (layoutInfo) {
-          preLayoutInfos.push([key, layoutInfo]);
-        }
-      }
-      anchor = this._scrollAnchor.captureBeforeLayout(anchorInfo, preLayoutInfos, this.visibleRect);
+      anchor = this._scrollAnchor.captureBeforeLayout(
+        anchorInfo,
+        this.getVisibleLayoutInfos(),
+        this.visibleRect
+      );
     }
 
     let previousContentSize = this.contentSize;
@@ -196,18 +200,29 @@ export class Virtualizer<T extends object, V> {
     let rawContentSize = this.layout.getContentSize();
     (this as Mutable<this>).contentSize = new Size(rawContentSize.width, rawContentSize.height);
 
+    // Decide whether the change that triggered this relayout was at the anchored edge. If items
+    // resized but none of them were the newest content, we don't follow the edge and we keep the
+    // user's reading position instead.
+    let changeIsAtEdge = !this._hadItemResize || this._batchIncludedNewestContent;
+
     let target = this._scrollAnchor.resolveAfterLayout({
       anchorInfo,
       anchor,
-      postLayoutInfos: anchorInfo ? this.getVisibleLayoutInfos() : new Map(),
       previousVisibleRect,
       previousContentSize,
       contentSize: this.contentSize,
       itemSizeChanged: context.itemSizeChanged ?? false,
       isScrolling: this._isScrolling,
-      getLayoutInfo: (key: Key) => this.layout.getLayoutInfo(key)
+      getLayoutInfo: (key: Key) => this.layout.getLayoutInfo(key),
+      changeIsAtEdge
     });
 
+    // Clear these flags because a relayout can also run for reasons unrelated to a resize
+    // (scrolling, a new message, a window resize). If we left the flags set, the next relayout
+    // would still see this pass's "an item resized / it was the newest" values and make the wrong call.
+    this._hadItemResize = false;
+    this._batchIncludedNewestContent = false;
+
     if (target) {
       // Queues a new render cycle. Return early to skip updateSubviews — running it now
       // would position views against the old visibleRect, causing a flash before the
@@ -447,9 +462,28 @@ export class Virtualizer<T extends object, V> {
 
     let changed = this.layout.updateItemSize(key, size);
     if (changed) {
+      this._hadItemResize = true;
+      // "Batch" refers to the set of updateItemSize calls that happen between one relayout and the next
+      this._batchIncludedNewestContent ||= this.isNewestContent(key);
       this.invalidate({
         itemSizeChanged: true
       });
     }
   }
+
+  /**
+   * Whether `key` is the newest real item in the collection — i.e. the last non-loader node.
+   */
+  private isNewestContent(key: Key): boolean {
+    let lastKey = this.collection.getLastKey();
+    while (lastKey != null) {
+      let node = this.collection.getItem(lastKey) as {type?: string} | null;
+      if (node?.type !== 'loader') {
+        break;
+      }
+      lastKey = this.collection.getKeyBefore(lastKey);
+    }
+
+    return lastKey != null && lastKey === key;
+  }
 }