fix: focus restoring in collections

packages/@react-spectrum/listbox/stories/ListBox.stories.tsx

@@ -20,9 +20,9 @@ import Book from '@spectrum-icons/workflow/Book';
 import Copy from '@spectrum-icons/workflow/Copy';
 import Cut from '@spectrum-icons/workflow/Cut';
 import Delete from '@spectrum-icons/workflow/Delete';
+import {DOMRefValue, Key} from '@react-types/shared';
 import {FocusScope} from '@react-aria/focus';
 import {Item, ListBox, Section} from '../';
-import {Key} from '@react-types/shared';
 import {Label} from '@react-spectrum/label';
 import {Meta, StoryFn, StoryObj} from '@storybook/react';
 import Paste from '@spectrum-icons/workflow/Paste';
@@ -946,7 +946,7 @@ export let FocusExample = (args: Omit<SpectrumListBoxProps<any>, 'children'> = {
   });
 
   let [dialog, setDialog] = useState<{action: Key} | null>(null);
-  let ref = useRef(null);
+  let ref = useRef<DOMRefValue<HTMLDivElement>>(null);
 
   return (
     <FocusScope>
@@ -989,7 +989,13 @@ export let FocusExample = (args: Omit<SpectrumListBoxProps<any>, 'children'> = {
               title="Delete"
               variant="destructive"
               primaryActionLabel="Delete"
-              onPrimaryAction={() => tree.removeSelectedItems()}>
+              onPrimaryAction={() => {
+                tree.removeSelectedItems();
+                // wait for inert to clear before focusing
+                requestAnimationFrame(() => {
+                  ref.current?.UNSAFE_getDOMNode()?.focus();
+                });
+              }}>
               Are you sure you want to delete {tree.selectedKeys.size === 1 ? '1 item' : `${tree.selectedKeys.size} items`}?
             </AlertDialog>
           }

packages/@react-spectrum/listbox/test/ListBox.test.js

@@ -1042,54 +1042,67 @@ describe('ListBox', function () {
       act(() => jest.runAllTimers());
       let listbox = tree.getByRole('listbox');
       let options = within(listbox).getAllByRole('option');
-      act(() => {options[2].focus();});
+      expect(options.length).toBe(6);
+      // Go to *Snake* and select it
+      await user.tab();
+      await user.keyboard('{ArrowDown}');
+      await user.keyboard('{ArrowDown}');
       expect(options[2]).toHaveAttribute('aria-posinset', '3');
       expect(options[2]).toHaveAttribute('aria-setsize', '6');
       expect(document.activeElement).toBe(options[2]);
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
+
+      // Remove *Snake*
       let removeButton = tree.getByRole('button');
       expect(removeButton).toBeInTheDocument();
-      act(() => {removeButton.focus();});
-      expect(document.activeElement).toBe(removeButton);
       await user.click(removeButton);
       act(() => jest.runAllTimers());
       let confirmationDialog = tree.getByRole('alertdialog');
       expect(document.activeElement).toBe(confirmationDialog);
       let confirmationDialogButton = within(confirmationDialog).getByRole('button');
-      expect(confirmationDialogButton).toBeInTheDocument();
       await user.click(confirmationDialogButton);
       act(() => jest.runAllTimers());
       options = within(listbox).getAllByRole('option');
       expect(options.length).toBe(5);
       act(() => jest.runAllTimers());
       expect(confirmationDialog).not.toBeInTheDocument();
+
+      // Dialog returns focus to the ListBox which forwards it to the options
       expect(document.activeElement).toBe(options[2]);
       expect(options[2]).toHaveAttribute('aria-posinset', '3');
       expect(options[2]).toHaveAttribute('aria-setsize', '5');
-      act(() => {options[1].focus();});
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+
+      // Select option
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
-      act(() => {options[0].focus();});
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+
+      // Go to option 0, *Aardvark* and select it too
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      expect(document.activeElement).toBe(options[0]);
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
-      act(() => {options[0].focus();});
+
+      // Remove the two selected items
       removeButton = tree.getByRole('button');
-      expect(removeButton).toBeInTheDocument();
-      act(() => {removeButton.focus();});
+      await user.tab({shift: true});
       expect(document.activeElement).toBe(removeButton);
       await user.click(removeButton);
       act(() => jest.runAllTimers());
       confirmationDialog = tree.getByRole('alertdialog');
       expect(document.activeElement).toBe(confirmationDialog);
       confirmationDialogButton = within(confirmationDialog).getByRole('button');
-      expect(confirmationDialogButton).toBeInTheDocument();
       await user.click(confirmationDialogButton);
       act(() => jest.runAllTimers());
       options = within(listbox).getAllByRole('option');
       expect(options.length).toBe(3);
       act(() => jest.runAllTimers());
       expect(confirmationDialog).not.toBeInTheDocument();
+
+      // Dialog returns focus to the ListBox which forwards it to the options
       expect(document.activeElement).toBe(options[0]);
       expect(options[0]).toHaveAttribute('aria-posinset', '1');
       expect(options[0]).toHaveAttribute('aria-setsize', '3');

packages/@react-stately/list/src/useListState.ts

@@ -90,25 +90,35 @@ function useFocusedKeyReset<T>(collection: Collection<Node<T>>, selectionManager
   useEffect(() => {
     if (selectionManager.focusedKey != null && !collection.getItem(selectionManager.focusedKey) && cachedCollection.current) {
       const startItem = cachedCollection.current.getItem(selectionManager.focusedKey);
-      const cachedItemNodes = [...cachedCollection.current.getKeys()].map(
-        key => {
-          const itemNode = cachedCollection.current!.getItem(key);
-          return itemNode?.type === 'item' ? itemNode : null;
-        }
-      ).filter(node => node !== null);
-      const itemNodes = [...collection.getKeys()].map(
-        key => {
-          const itemNode = collection.getItem(key);
-          return itemNode?.type === 'item' ? itemNode : null;
+
+      // Helper to get all item nodes from a collection (flattening sections)
+      const getAllItemNodes = (coll: Collection<Node<T>>): Node<T>[] => {
+        const items: Node<T>[] = [];
+        for (let node of coll) {
+          if (node.type === 'item') {
+            items.push(node);
+          } else if (node.type === 'section' && coll.getChildren?.(node.key)) {
+            for (let child of coll.getChildren(node.key)) {
+              if (child.type === 'item') {
+                items.push(child);
+              }
+            }
+          }
         }
-      ).filter(node => node !== null);
-      const diff: number = (cachedItemNodes?.length ?? 0) - (itemNodes?.length ?? 0);
+        return items;
+      };
+
+      const cachedItemNodes = getAllItemNodes(cachedCollection.current);
+      const itemNodes = getAllItemNodes(collection);
+
+      // Count how many items were removed before the focused item's original index
+      const itemNodesKeys = new Set(itemNodes.map(node => node.key));
+      const removedBeforeCount = cachedItemNodes.filter((node, idx) =>
+        idx < (startItem?.index ?? 0) && !itemNodesKeys.has(node.key)
+      ).length;
+
       let index = Math.min(
-        (
-          diff > 1 ?
-          Math.max((startItem?.index ?? 0) - diff + 1, 0) :
-          startItem?.index ?? 0
-        ),
+        Math.max((startItem?.index ?? 0) - removedBeforeCount, 0),
         (itemNodes?.length ?? 0) - 1);
       let newNode: Node<T> | null = null;
       let isReverseSearching = false;

packages/react-aria-components/test/TagGroup.test.js

@@ -12,7 +12,7 @@
 
 import {act, fireEvent, mockClickDefault, pointerMap, render} from '@react-spectrum/test-utils-internal';
 import {Button, Label, RouterProvider, Tag, TagGroup, TagList, Text, Tooltip, TooltipTrigger} from '../';
-import React from 'react';
+import React, {useRef} from 'react';
 import {useListData} from '@react-stately/data';
 import {User} from '@react-aria/test-utils';
 import userEvent from '@testing-library/user-event';
@@ -549,6 +549,77 @@ describe('TagGroup', () => {
     expect(onRemove).toHaveBeenLastCalledWith(new Set(['dog']));
   });
 
+  it('should maintain item order when adding new items', async () => {
+    function MyTag(props) {
+      return (
+        <Tag
+          {...props}
+          style={({isSelected}) => ({border: '1px solid gray', borderRadius: 4, padding: '0 4px', background: isSelected ? 'black' : '', color: isSelected ? 'white' : '', cursor: props.href ? 'pointer' : 'default'})} />
+      );
+    }
+    function Example() {
+      const list = useListData({
+        initialItems: []
+      });
+
+      const nextIdRef = useRef(0);
+
+      const insertItem = () => {
+        const id = nextIdRef.current++;
+        list.insert(0, {
+          id,
+          label: `Item ${id + 1}`
+        });
+      };
+
+      return (
+        <div>
+          <Button onPress={insertItem}>Insert item</Button>
+          <TagGroup onRemove={keys => list.remove(...keys)}>
+            <Label>Categories</Label>
+            <TagList style={{display: 'flex', gap: 4}} items={list.items} renderEmptyState={() => 'No categories.'}>
+              {item => <MyTag textValue={item.label}>{item.label}<Button slot="remove">X</Button></MyTag>}
+            </TagList>
+          </TagGroup>
+        </div>
+      );
+    }
+    let {getAllByRole, queryAllByRole, getByRole} = render(<Example />);
+    let addButton = getAllByRole('button')[0];
+    let tagGroup = getByRole('group');
+
+    await user.click(addButton);
+    await user.click(addButton);
+    await user.click(addButton);
+    await user.click(addButton);
+    act(() => jest.runAllTimers());
+    let items = getAllByRole('row');
+    expect(items).toHaveLength(4);
+    expect(items[0]).toHaveTextContent('Item 4');
+
+    await user.tab();
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(3);
+    expect(items[0]).toHaveTextContent('Item 3');
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(2);
+    expect(items[0]).toHaveTextContent('Item 2');
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(1);
+    expect(items[0]).toHaveTextContent('Item 1');
+
+    await user.keyboard('{Delete}');
+    let noItems = queryAllByRole('row');
+    expect(noItems).toHaveLength(0);
+    expect(document.activeElement).toBe(tagGroup);
+  });
+
   describe('shouldSelectOnPressUp', () => {
     it('should select an item on pressing down when shouldSelectOnPressUp is not provided', async () => {
       let onSelectionChange = jest.fn();
@@ -588,7 +659,7 @@ describe('TagGroup', () => {
   });
 
   describe('press events', () => {
-    it.only.each`
+    it.each`
       interactionType
       ${'mouse'}
       ${'keyboard'}