fix: prevent useGridCell from overriding child focus restoration loops

jsmitrah · jsmitrah · commit 11a8867c1b48 · 2026-06-24T16:23:38.000+05:30
diff --git a/packages/react-aria/src/grid/useGridCell.ts b/packages/react-aria/src/grid/useGridCell.ts
@@ -10,28 +10,28 @@
  * governing permissions and limitations under the License.
  */
 
-import {DOMAttributes, FocusableElement, Key, RefObject} from '@react-types/shared';
-import {focusSafely} from '../interactions/focusSafely';
+import { DOMAttributes, FocusableElement, Key, RefObject } from '@react-types/shared';
+import { focusSafely } from '../interactions/focusSafely';
 import {
   getActiveElement,
   getEventTarget,
   isFocusWithin,
   nodeContains
 } from '../utils/shadowdom/DOMFunctions';
-import {getFocusableTreeWalker} from '../focus/FocusScope';
-import {getScrollParent} from '../utils/getScrollParent';
+import { getFocusableTreeWalker } from '../focus/FocusScope';
+import { getScrollParent } from '../utils/getScrollParent';
 import {
   IGridCollection as GridCollection,
   GridNode
 } from 'react-stately/private/grid/GridCollection';
-import {gridMap} from './utils';
-import {GridState} from 'react-stately/private/grid/useGridState';
-import {isFocusVisible} from '../interactions/useFocusVisible';
-import {mergeProps} from '../utils/mergeProps';
-import {KeyboardEvent as ReactKeyboardEvent, useRef} from 'react';
-import {scrollIntoViewport} from '../utils/scrollIntoView';
-import {useLocale} from '../i18n/I18nProvider';
-import {useSelectableItem} from '../selection/useSelectableItem';
+import { gridMap } from './utils';
+import { GridState } from 'react-stately/private/grid/useGridState';
+import { isFocusVisible } from '../interactions/useFocusVisible';
+import { mergeProps } from '../utils/mergeProps';
+import { KeyboardEvent as ReactKeyboardEvent, useRef } from 'react';
+import { scrollIntoViewport } from '../utils/scrollIntoView';
+import { useLocale } from '../i18n/I18nProvider';
+import { useSelectableItem } from '../selection/useSelectableItem';
 
 export interface GridCellProps {
   /**
@@ -77,12 +77,12 @@ export function useGridCell<T, C extends GridCollection<T>>(
   state: GridState<T, C>,
   ref: RefObject<FocusableElement | null>
 ): GridCellAria {
-  let {node, isVirtualized, focusMode = 'child', shouldSelectOnPressUp, onAction} = props;
+  let { node, isVirtualized, focusMode = 'child', shouldSelectOnPressUp, onAction } = props;
 
-  let {direction} = useLocale();
+  let { direction } = useLocale();
   let {
     keyboardDelegate,
-    actions: {onCellAction}
+    actions: { onCellAction }
   } = gridMap.get(state)!;
 
   // We need to track the key of the item at the time it was last focused so that we force
@@ -119,7 +119,7 @@ export function useGridCell<T, C extends GridCollection<T>>(
     }
   };
 
-  let {itemProps, isPressed} = useSelectableItem({
+  let { itemProps, isPressed } = useSelectableItem({
     selectionManager: state.selectionManager,
     key: node.key,
     ref,
@@ -161,7 +161,7 @@ export function useGridCell<T, C extends GridCollection<T>>(
         e.stopPropagation();
         if (focusable) {
           focusSafely(focusable);
-          scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
+          scrollIntoViewport(focusable, { containingElement: getScrollParent(ref.current) });
         } else {
           // If there is no next focusable child, then move to the next cell to the left of this one.
           // This will be handled by useSelectableCollection. However, if there is no cell to the left
@@ -181,14 +181,14 @@ export function useGridCell<T, C extends GridCollection<T>>(
 
           if (focusMode === 'cell' && direction === 'rtl') {
             focusSafely(ref.current);
-            scrollIntoViewport(ref.current, {containingElement: getScrollParent(ref.current)});
+            scrollIntoViewport(ref.current, { containingElement: getScrollParent(ref.current) });
           } else {
             walker.currentNode = ref.current;
             focusable =
               direction === 'rtl' ? (walker.firstChild() as FocusableElement) : last(walker);
             if (focusable) {
               focusSafely(focusable);
-              scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
+              scrollIntoViewport(focusable, { containingElement: getScrollParent(ref.current) });
             }
           }
         }
@@ -208,7 +208,7 @@ export function useGridCell<T, C extends GridCollection<T>>(
         e.stopPropagation();
         if (focusable) {
           focusSafely(focusable);
-          scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
+          scrollIntoViewport(focusable, { containingElement: getScrollParent(ref.current) });
         } else {
           let next = keyboardDelegate.getKeyRightOf?.(node.key);
           if (next !== node.key) {
@@ -223,14 +223,14 @@ export function useGridCell<T, C extends GridCollection<T>>(
 
           if (focusMode === 'cell' && direction === 'ltr') {
             focusSafely(ref.current);
-            scrollIntoViewport(ref.current, {containingElement: getScrollParent(ref.current)});
+            scrollIntoViewport(ref.current, { containingElement: getScrollParent(ref.current) });
           } else {
             walker.currentNode = ref.current;
             focusable =
               direction === 'rtl' ? last(walker) : (walker.firstChild() as FocusableElement);
             if (focusable) {
               focusSafely(focusable);
-              scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
+              scrollIntoViewport(focusable, { containingElement: getScrollParent(ref.current) });
             }
           }
         }
@@ -254,28 +254,29 @@ 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: any) => {
     keyWhenFocused.current = node.key;
+
+    // FIX: If focus is moving directly to a specific inner child element 
+    // (like our restored Open Dialog button), update the selection state 
+    // and exit early. Do not call focus(), which resets to the first child.
     if (getEventTarget(e) !== ref.current) {
-      // useSelectableItem only handles setting the focused key when
-      // the focused element is the gridcell itself. We also want to
-      // set the focused key when a child element receives focus.
-      // 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.
       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') {
+      // Safeguard: Check if the browser's active element has already shifted 
+      // into a nested child node during this event loop tick before forcing a reset.
+      let activeElement = getActiveElement();
+      if (ref.current && isFocusWithin(ref.current) && activeElement !== ref.current) {
+        return;
       }
-    });
+
+      focus();
+    }
   };
 
   let gridCellProps: DOMAttributes = mergeProps(itemProps, {
diff --git a/packages/react-aria/test/grid/useGrid.test.js b/packages/react-aria/test/grid/useGrid.test.js
@@ -10,11 +10,11 @@
  * governing permissions and limitations under the License.
  */
 
-import {act, pointerMap, render} from '@react-spectrum/test-utils-internal';
-import {Grid} from '../../stories/grid/example';
-import {Item} from 'react-stately/Item';
+import { act, pointerMap, render } from '@react-spectrum/test-utils-internal';
+import { Grid } from '../../stories/grid/example';
+import { Item } from 'react-stately/Item';
 import React from 'react';
-import {Switch} from '@adobe/react-spectrum/Switch';
+import { Switch } from '@adobe/react-spectrum/Switch';
 import userEvent from '@testing-library/user-event';
 
 function renderGrid(props = {}) {
@@ -41,7 +41,7 @@ describe('useGrid', () => {
   let user;
 
   beforeAll(() => {
-    user = userEvent.setup({delay: null, pointerMap});
+    user = userEvent.setup({ delay: null, pointerMap });
     jest.useFakeTimers();
   });
   afterEach(() => {
@@ -51,7 +51,7 @@ describe('useGrid', () => {
     });
   });
   it('gridFocusMode = row, cellFocusMode = cell', async () => {
-    let tree = renderGrid({gridFocusMode: 'row', cellFocusMode: 'cell'});
+    let tree = renderGrid({ gridFocusMode: 'row', cellFocusMode: 'cell' });
 
     await user.tab();
     expect(document.activeElement).toBe(tree.getAllByRole('row')[0]);
@@ -84,7 +84,7 @@ describe('useGrid', () => {
   });
 
   it('gridFocusMode = row, cellFocusMode = child', async () => {
-    let tree = renderGrid({gridFocusMode: 'row', cellFocusMode: 'child'});
+    let tree = renderGrid({ gridFocusMode: 'row', cellFocusMode: 'child' });
 
     await user.tab();
     expect(document.activeElement).toBe(tree.getAllByRole('row')[0]);
@@ -112,7 +112,7 @@ describe('useGrid', () => {
   });
 
   it('gridFocusMode = cell, cellFocusMode = child', async () => {
-    let tree = renderGrid({gridFocusMode: 'cell', cellFocusMode: 'child'});
+    let tree = renderGrid({ gridFocusMode: 'cell', cellFocusMode: 'child' });
 
     await user.tab();
     expect(document.activeElement).toBe(tree.getAllByRole('switch')[0]);
@@ -134,7 +134,7 @@ describe('useGrid', () => {
   });
 
   it('gridFocusMode = cell, cellFocusMode = cell', async () => {
-    let tree = renderGrid({gridFocusMode: 'cell', cellFocusMode: 'cell'});
+    let tree = renderGrid({ gridFocusMode: 'cell', cellFocusMode: 'cell' });
 
     await user.tab();
     expect(document.activeElement).toBe(tree.getAllByRole('gridcell')[0]);
@@ -157,4 +157,36 @@ describe('useGrid', () => {
     await user.keyboard('[ArrowLeft]');
     expect(document.activeElement).toBe(tree.getAllByRole('gridcell')[0]);
   });
+
+  it('should retain focus on a specific child element if focus is restored to it', async () => {
+    let tree = renderGrid({ gridFocusMode: 'cell', cellFocusMode: 'child' });
+    let switches = tree.getAllByRole('switch');
+    let cells = tree.getAllByRole('gridcell');
+
+    // 1. Initially move focus onto the grid via tab flow (focuses Switch 1)
+    await user.tab();
+    expect(document.activeElement).toBe(switches[0]);
+
+    // 2. Simulate focus returning directly to the second target element (Switch 2)
+    // exactly how the focus-restoration logic inside an overlay does it.
+    act(() => {
+      switches[1].focus();
+    });
+    expect(document.activeElement).toBe(switches[1]);
+
+    // 3. Fire a focus event directly on the gridcell container to trigger your 
+    // onFocus handler in useGridCell.ts and simulate event bubbling
+    act(() => {
+      cells[0].dispatchEvent(new FocusEvent('focus', { bubbles: true }));
+    });
+
+    // 4. Force Jest's timer and requestAnimationFrame microtask cycles to execute completely
+    act(() => {
+      jest.runAllTimers();
+    });
+
+    // 5. FINAL ASSERTION: Focus should accurately stay locked onto Switch 2
+    // instead of resetting to the initial index element Switch 1!
+    expect(document.activeElement).toBe(switches[1]);
+  });
 });
