fix: #8686 prevent useGridCell from overriding child focus restoration loops

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

@@ -315,6 +315,10 @@ let DynamicTable = ({tableProps, tableHeaderProps, tableBodyProps, rowProps}) =>
 
 let renderTable = props => render(<TestTable {...props} />);
 
+if (typeof Symbol.dispose === 'undefined') {
+  Object.defineProperty(Symbol, 'dispose', {value: Symbol('dispose')});
+}
+
 describe('Table', () => {
   let user;
   let testUtilUser = new User();
@@ -3535,6 +3539,43 @@ describe('Table', () => {
     expect(tableTester.getFooterRows()).toHaveLength(1);
     expect(tableTester.getFooterRows()[0]).toHaveTextContent('Blah');
   });
+
+  it('should return focus to the specific button that opened a modal inside a cell, not the first button', async () => {
+    let {getByRole} = render(
+      <Table aria-label="Books">
+        <TableHeader>
+          <Column isRowHeader>Title</Column>
+          <Column>Actions</Column>
+        </TableHeader>
+        <TableBody>
+          <Row>
+            <Cell>The Great Gatsby</Cell>
+            <Cell>
+              <Button data-testid="first-btn">Edit</Button>
+              <Button data-testid="trigger-btn">Open Dialog</Button>
+            </Cell>
+          </Row>
+        </TableBody>
+      </Table>
+    );
+
+    let triggerButton = getByRole('button', {name: 'Open Dialog'});
+    let firstButton = getByRole('button', {name: 'Edit'});
+    let actionsCell = getByRole('gridcell');
+
+    await user.tab();
+    await user.keyboard('[ArrowRight]');
+    await user.keyboard('[ArrowRight]');
+    await user.keyboard('[ArrowRight]');
+    expect(document.activeElement).toBe(triggerButton);
+
+    act(() => {
+      actionsCell.focus();
+    });
+
+    expect(document.activeElement).toBe(triggerButton);
+    expect(document.activeElement).not.toBe(firstButton);
+  });
 });
 
 function HidingColumnsExample({dynamic = false}) {

packages/react-aria/src/grid/useGridCell.ts

@@ -89,17 +89,38 @@ export function useGridCell<T, C extends GridCollection<T>>(
   // focus to go to the item when the DOM node is reused for a different item in a virtualizer.
   let keyWhenFocused = useRef<Key | null>(null);
 
+  // Tracks the specific focusable child that was last focused within this cell.
+  let lastFocusedChild = useRef<FocusableElement | null>(null);
+
   // Handles focusing the cell. If there is a focusable child,
   // it is focused, otherwise the cell itself is focused.
   let focus = () => {
     if (ref.current) {
       let treeWalker = getFocusableTreeWalker(ref.current);
       if (focusMode === 'child') {
+        let activeElement = getActiveElement();
+
         // If focus is already on a focusable child within the cell, early return so we don't shift focus
-        if (isFocusWithin(ref.current) && ref.current !== getActiveElement()) {
+        if (isFocusWithin(ref.current) && ref.current !== activeElement) {
           return;
         }
 
+        // Focus counts as "disrupted" (as opposed to a fresh keyboard entry from a
+        // sibling row/cell) when it's on the cell itself or fell back to body/nothing,
+        // e.g. because an overlay closed and removed the element that had focus. In
+        // that case, restore the child that was last focused instead of defaulting to
+        // childFocusStrategy's first/last child.
+        let ownerDocument = ref.current.ownerDocument;
+        let isDisrupted =
+          ref.current === activeElement || !activeElement || activeElement === ownerDocument.body;
+        if (isDisrupted) {
+          let lastChild = lastFocusedChild.current;
+          if (lastChild && nodeContains(ref.current, lastChild)) {
+            focusSafely(lastChild);
+            return;
+          }
+        }
+
         let focusable =
           state.selectionManager.childFocusStrategy === 'last'
             ? last(treeWalker)
@@ -254,7 +275,7 @@ export function useGridCell<T, C extends GridCollection<T>>(
 
   // Grid cells can have focusable elements inside them. In this case, focus should
   // be marshalled to that element rather than focusing the cell itself.
-  let onFocus = e => {
+  let onFocus = (e: FocusEvent) => {
     keyWhenFocused.current = node.key;
     if (getEventTarget(e) !== ref.current) {
       // useSelectableItem only handles setting the focused key when
@@ -263,19 +284,28 @@ export function useGridCell<T, C extends GridCollection<T>>(
       // If focus is currently visible (e.g. the user is navigating with the keyboard),
       // then skip this. We want to restore focus to the previously focused row/cell
       // in that case since the table should act like a single tab stop.
+
+      // Remember which child was focused so that if something later forces focus
+      // back to this cell (e.g. a closing overlay), focus() can restore it directly
+      // instead of falling back to the first/last focusable child. Only do this for
+      // actual DOM descendants of the cell -- portalled content (e.g. a dialog opened
+      // from within the cell) can still reach this handler because React dispatches
+      // events along the component tree rather than the DOM tree for portals, even
+      // though it isn't really inside this cell in the DOM.
+      let target = getEventTarget(e) as FocusableElement;
+      if (ref.current && nodeContains(ref.current, target)) {
+        lastFocusedChild.current = target;
+      }
+
       if (!isFocusVisible()) {
         state.selectionManager.setFocusedKey(node.key);
       }
       return;
     }
 
-    // If the cell itself is focused, wait a frame so that focus finishes propagatating
-    // up to the tree, and move focus to a focusable child if possible.
-    requestAnimationFrame(() => {
-      if (focusMode === 'child' && getActiveElement() === ref.current) {
-        focus();
-      }
-    });
+    if (focusMode === 'child') {
+      focus();
+    }
   };
 
   // oxlint-disable-next-line react/react-compiler

packages/react-aria/test/grid/useGrid.test.js

@@ -157,4 +157,22 @@ describe('useGrid', () => {
     await user.keyboard('[ArrowLeft]');
     expect(document.activeElement).toBe(tree.getAllByRole('gridcell')[0]);
   });
+
+  it('should restore focus to the child that was last focused within a cell, not the first child', async () => {
+    let tree = renderGrid({gridFocusMode: 'cell', cellFocusMode: 'child'});
+    let switches = tree.getAllByRole('switch');
+    let cells = tree.getAllByRole('gridcell');
+
+    await user.tab();
+    expect(document.activeElement).toBe(switches[0]);
+
+    await user.keyboard('[ArrowRight]');
+    expect(document.activeElement).toBe(switches[1]);
+
+    act(() => {
+      cells[0].focus();
+    });
+
+    expect(document.activeElement).toBe(switches[1]);
+  });
 });