Fix Select + Dialog focus issue (#1608)

Author: mdodgelooker

packages/components/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Fixed
+
+- `Select` inside a `Dialog` losing focus after clicking an option
+
 ### Changed
 
 - `InputFiltersChipEditor` can use CheckboxGroup or RadioGroup.

packages/components/src/utils/useFocusTrap.ts

@@ -74,10 +74,20 @@ export function useFocusTrap(
           clickOutsideDeactivates: true,
           escapeDeactivates: false,
           fallbackFocus: element,
+          initialFocus:
+            autoFocusElement || surfaceElement
+              ? () => {
+                  // Check first if focus is already within the element
+                  // otherwise if focus trap is disabled then re-enabled
+                  // (e.g. when a Select closes inside a Dialog)
+                  // focus will move unnecessarily back to the surface / autoFocus element
+                  if (element.contains(document.activeElement)) {
+                    return document.activeElement as HTMLElement
+                  }
+                  return autoFocusElement || surfaceElement
+                }
+              : undefined,
           onDeactivate: () => setOff(),
-          ...(autoFocusElement || surfaceElement
-            ? { initialFocus: autoFocusElement || surfaceElement }
-            : {}),
         })
       }
       disableFocusTrap && disableFocusTrap()