fix(compass-components): Keep the focus in the grid item that received the event

packages/compass-components/src/components/virtual-grid.tsx

@@ -9,10 +9,7 @@ import React, {
 import { css, cx } from '@leafygreen-ui/emotion';
 import { FixedSizeList } from 'react-window';
 import { useDOMRect } from '../hooks/use-dom-rect';
-import {
-  useVirtualGridArrowNavigation,
-  useVirtualRovingTabIndex,
-} from '../hooks/use-virtual-grid';
+import { useVirtualGridArrowNavigation } from '../hooks/use-virtual-grid';
 import { mergeProps } from '../utils/merge-props';
 
 type RenderItem = React.FunctionComponent<
@@ -49,6 +46,10 @@ type VirtualGridProps = {
    * correctly
    */
   renderItem: RenderItem;
+  /**
+   * Custom grid item key (default is item index)
+   */
+  itemKey?: (index: number) => React.Key | null | undefined;
   /**
    * Header content height
    */
@@ -77,7 +78,11 @@ type VirtualGridProps = {
     cell?: string;
   };
 
-  itemKey?: (index: number) => React.Key | null | undefined;
+  /**
+   * Set to `false` of you want the last focused item to be preserved between
+   * focus / blur (default: true)
+   */
+  resetActiveItemOnBlur?: boolean;
 };
 
 const GridContext = createContext<
@@ -130,10 +135,10 @@ const GridWithHeader = forwardRef<
       }}
       {...props}
     >
-      <div className={classNames?.header}>
+      <div style={{ height: headerHeight }} className={classNames?.header}>
         {React.createElement(renderHeader, {})}
       </div>
-      <div {...gridProps}>
+      <div style={{ height: style.height }} {...gridProps}>
         {itemsCount === 0 && renderEmptyList
           ? React.createElement(renderEmptyList, {})
           : children}
@@ -245,6 +250,7 @@ export const VirtualGrid = forwardRef<
     overscanCount = 3,
     classNames,
     itemKey,
+    resetActiveItemOnBlur,
     ...containerProps
   },
   ref
@@ -269,13 +275,10 @@ export const VirtualGrid = forwardRef<
       itemsCount,
       colCount,
       rowCount,
+      onFocusMove,
+      resetActiveItemOnBlur,
     });
 
-  const rovingFocusProps = useVirtualRovingTabIndex<HTMLDivElement>({
-    currentTabbable,
-    onFocusMove,
-  });
-
   const gridContainerProps = mergeProps(
     { ref, className: cx(container, classNames?.container) },
     containerProps,
@@ -299,8 +302,7 @@ export const VirtualGrid = forwardRef<
             'aria-rowcount': rowCount,
             className: cx(grid, classNames?.grid),
           },
-          navigationProps,
-          rovingFocusProps
+          navigationProps
         ),
         itemKey,
         renderEmptyList,

packages/compass-components/src/hooks/use-default-action.ts

@@ -11,14 +11,6 @@ import { useCallback } from 'react';
 export function useDefaultAction<T>(
   onDefaultAction: (evt: React.KeyboardEvent<T> | React.MouseEvent<T>) => void
 ): React.HTMLAttributes<T> {
-  // Prevent event from possibly causing bubbled focus on parent element, if
-  // something is interacting with this component using mouse, we want to
-  // prevent anything from bubbling
-  const onMouseDown = useCallback((evt: React.MouseEvent<T>) => {
-    evt.preventDefault();
-    evt.stopPropagation();
-  }, []);
-
   const onClick = useCallback(
     (evt: React.MouseEvent<T>) => {
       evt.stopPropagation();
@@ -42,5 +34,5 @@ export function useDefaultAction<T>(
     [onDefaultAction]
   );
 
-  return { onMouseDown, onClick, onKeyDown };
+  return { onClick, onKeyDown };
 }

packages/compass-components/src/hooks/use-virtual-grid.test.tsx

@@ -1,13 +1,8 @@
-/* eslint-disable react/prop-types */
 import React from 'react';
 import { render, screen, cleanup } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { expect } from 'chai';
-import { mergeProps } from '../utils/merge-props';
-import {
-  useVirtualGridArrowNavigation,
-  useVirtualRovingTabIndex,
-} from './use-virtual-grid';
+import { useVirtualGridArrowNavigation } from './use-virtual-grid';
 
 const TestGrid: React.FunctionComponent<{
   rowCount?: number;
@@ -28,17 +23,11 @@ const TestGrid: React.FunctionComponent<{
       colCount,
       itemsCount: rowCount * colCount,
       defaultCurrentTabbable,
+      onFocusMove,
     });
 
-  const rovingTabIndexProps = useVirtualRovingTabIndex<HTMLDivElement>({
-    currentTabbable,
-    onFocusMove,
-  });
-
-  const props = mergeProps(arrowNavigationProps, rovingTabIndexProps);
-
   return (
-    <div role="grid" aria-rowcount={rowCount} {...props}>
+    <div role="grid" aria-rowcount={rowCount} {...arrowNavigationProps}>
       {Array.from({ length: rowCount }, (_, row) => (
         <div key={row} role="row" aria-rowindex={row + 1}>
           {Array.from({ length: colCount }, (_, col) => {
@@ -258,4 +247,11 @@ describe('virtual grid keyboard navigation', function () {
       expect(screen.getByText('5-3')).to.eq(document.activeElement);
     });
   });
+
+  it('should keep focus on the element that was interacted with', function () {
+    render(<TestGrid></TestGrid>);
+    expect(document.body).to.eq(document.activeElement);
+    userEvent.click(screen.getByText('2-3'));
+    expect(screen.getByText('2-3')).to.eq(document.activeElement);
+  });
 });

packages/compass-components/src/hooks/use-virtual-grid.ts

@@ -1,6 +1,43 @@
+import type React from 'react';
 import { useRef, useState, useEffect, useCallback } from 'react';
-import { useFocusState, FocusState } from './use-focus-hover';
 
+function closest(
+  node: HTMLElement | null,
+  cond: ((node: HTMLElement) => boolean) | string
+): HTMLElement | null {
+  if (typeof cond === 'string') {
+    return node?.closest(cond) ?? null;
+  }
+
+  let parent: HTMLElement | null = node;
+  while (parent) {
+    if (cond(parent)) {
+      return parent;
+    }
+    parent = parent.parentElement;
+  }
+  return null;
+}
+
+function vgridItemSelector(idx?: number): string {
+  return idx ? `[data-vlist-item-idx="${idx}"]` : '[data-vlist-item-idx]';
+}
+
+function getItemIndex(node: HTMLElement): number {
+  if (!node.dataset.vlistItemIdx) {
+    throw new Error('Trying to get vgrid item index from an non-item element');
+  }
+  return Number(node.dataset.vlistItemIdx);
+}
+
+/**
+ * Hook that adds support for the grid keyboard navigation while handling the
+ * focus using the roving tab index
+ *
+ * {@link https://www.w3.org/TR/wai-aria-1.1/#grid}
+ * {@link https://www.w3.org/TR/wai-aria-1.1/#gridcell}
+ * {@link https://www.w3.org/TR/wai-aria-practices-1.1/#kbd_roving_tabindex}
+ */
 export function useVirtualGridArrowNavigation<
   T extends HTMLElement = HTMLElement
 >({
@@ -10,25 +47,78 @@ export function useVirtualGridArrowNavigation<
   resetActiveItemOnBlur = true,
   pageSize = 3,
   defaultCurrentTabbable = 0,
+  onFocusMove,
 }: {
   colCount: number;
   rowCount: number;
   itemsCount: number;
   resetActiveItemOnBlur?: boolean;
   pageSize?: number;
   defaultCurrentTabbable?: number;
+  onFocusMove(idx: number): void;
 }): [React.HTMLProps<T>, number] {
   const rootNode = useRef<T | null>(null);
-  const [focusProps, focusState] = useFocusState();
+  const [tabIndex, setTabIndex] = useState<0 | -1>(0);
   const [currentTabbable, setCurrentTabbable] = useState(
     defaultCurrentTabbable
   );
 
-  useEffect(() => {
-    if (resetActiveItemOnBlur && focusState === FocusState.NoFocus) {
-      setCurrentTabbable(defaultCurrentTabbable);
+  const onFocus = useCallback(
+    (evt: React.FocusEvent) => {
+      // If we received focus on the grid container itself, this is a keyboard
+      // navigation, disable focus on the container to trigger a focus effect
+      // for the currentTabbable element
+      if (evt.target === evt.currentTarget) {
+        setTabIndex(-1);
+      } else {
+        const focusedItem = closest(
+          evt.target as HTMLElement,
+          vgridItemSelector()
+        );
+
+        // If focus was received somewhere inside grid item, disable focus on
+        // the container and mark item that got the interaction as the
+        // `currentTabbable` item
+        if (focusedItem) {
+          setTabIndex(-1);
+          setCurrentTabbable(getItemIndex(focusedItem));
+        }
+      }
+    },
+    [defaultCurrentTabbable]
+  );
+
+  const onBlur = useCallback(() => {
+    const isFocusInside =
+      closest(
+        document.activeElement as HTMLElement,
+        (node) => node === rootNode.current
+      ) !== null;
+
+    // If focus is outside of the grid container, make the whole container
+    // focusable again and reset tabbable item if needed
+    if (!isFocusInside) {
+      setTabIndex(0);
+      if (resetActiveItemOnBlur) {
+        setCurrentTabbable(defaultCurrentTabbable);
+      }
     }
-  }, [resetActiveItemOnBlur, focusState, defaultCurrentTabbable]);
+  }, [resetActiveItemOnBlur, defaultCurrentTabbable]);
+
+  const onMouseDown = useCallback((evt: React.MouseEvent) => {
+    const gridItem = closest(evt.target as HTMLElement, vgridItemSelector());
+    // If mousedown didn't originate in one of the grid items (we just clicked
+    // some empty space in the grid container), prevent default behavior to stop
+    // focus on the grid container from happening
+    if (!gridItem) {
+      evt.preventDefault();
+      // Simulate active element blur that normally happens when clicking a
+      // non-focusable element
+      (document.activeElement as HTMLElement)?.blur();
+    }
+  }, []);
+
+  const focusProps = { tabIndex, onFocus, onBlur, onMouseDown };
 
   const onKeyDown = useCallback(
     (evt: React.KeyboardEvent<T>) => {
@@ -116,53 +206,38 @@ export function useVirtualGridArrowNavigation<
     [currentTabbable, itemsCount, rowCount, colCount, pageSize]
   );
 
-  return [{ ref: rootNode, onKeyDown, ...focusProps }, currentTabbable];
-}
-
-export function useVirtualRovingTabIndex<T extends HTMLElement = HTMLElement>({
-  currentTabbable,
-  onFocusMove,
-}: {
-  currentTabbable: number;
-  onFocusMove(idx: number): void;
-}): React.HTMLProps<T> {
-  const rootNode = useRef<T | null>(null);
-  // We will set tabIndex on the parent element so that it can catch focus even
-  // if the currentTabbable is not rendered
-  const [tabIndex, setTabIndex] = useState<0 | -1>(0);
-  const [focusProps, focusState] = useFocusState();
-
-  // Focuses vlist item by id or falls back to the first focusable element in
-  // the container
-  const focusTabbable = useCallback(() => {
-    const selector =
-      currentTabbable >= 0
-        ? `[data-vlist-item-idx="${currentTabbable}"]`
-        : '[tabindex=0]';
-    rootNode.current?.querySelector<T>(selector)?.focus();
-  }, [rootNode, currentTabbable]);
+  const activeCurrentTabbable = tabIndex === 0 ? -1 : currentTabbable;
 
   useEffect(() => {
-    if (
-      [
-        FocusState.Focus,
-        FocusState.FocusVisible,
-        FocusState.FocusWithin,
-        FocusState.FocusWithinVisible,
-      ].includes(focusState)
-    ) {
-      setTabIndex(-1);
-      onFocusMove(currentTabbable);
-      const frame = requestAnimationFrame(() => {
-        focusTabbable();
-      });
-      return () => {
-        cancelAnimationFrame(frame);
-      };
-    } else {
-      setTabIndex(0);
+    // If we have an active current tabbable item (there is a focus somewhere in
+    // the grid container) ...
+    if (activeCurrentTabbable >= 0) {
+      const gridItem = closest(
+        document.activeElement as HTMLElement,
+        vgridItemSelector()
+      );
+      const shouldMoveFocus =
+        !gridItem || getItemIndex(gridItem) !== activeCurrentTabbable;
+
+      // ... and this item is currently not focused ...
+      if (shouldMoveFocus) {
+        // ... communicate that there will be a focus change happening (this is
+        // needed so that we can scroll invisible virtual item into view if
+        // needed) ...
+        onFocusMove(activeCurrentTabbable);
+        // ... and trigger a focus on the element after a frame delay, so that
+        // the item has time to scroll into view and render if needed
+        const frameId = requestAnimationFrame(() => {
+          rootNode.current
+            ?.querySelector<HTMLElement>(vgridItemSelector(currentTabbable))
+            ?.focus();
+        });
+        return () => {
+          cancelAnimationFrame(frameId);
+        };
+      }
     }
-  }, [focusState, onFocusMove, focusTabbable, currentTabbable]);
+  }, [activeCurrentTabbable]);
 
-  return { ref: rootNode, tabIndex, ...focusProps };
+  return [{ ref: rootNode, onKeyDown, ...focusProps }, activeCurrentTabbable];
 }

packages/compass-saved-aggregations-queries/src/components/aggregations-queries-list.tsx

@@ -63,7 +63,7 @@ const sortBy: { name: keyof Item; label: string }[] = [
 ];
 
 const headerStyles = css({
-  margin: spacing[3],
+  padding: spacing[3],
   display: 'flex',
   justifyContent: 'space-between',
 });
@@ -221,6 +221,7 @@ const AggregationsQueriesList = ({
         headerHeight={spacing[5] + 36}
         renderEmptyList={NoSearchResults}
         classNames={{ row: rowStyles }}
+        resetActiveItemOnBlur={false}
       ></VirtualGrid>
       <OpenItemModal></OpenItemModal>
       <EditItemModal></EditItemModal>

packages/compass-saved-aggregations-queries/src/components/saved-item-card.test.tsx

@@ -55,7 +55,6 @@ describe('SavedItemCard', function () {
     );
 
     userEvent.click(screen.getByText('My Awesome Query'));
-    userEvent.tab();
     userEvent.keyboard('{space}');
     userEvent.keyboard('{enter}');