[ESR-2405]-Restore focus to trigger button when pressing Escape

.changeset/fix-menu-escape-focus.md

@@ -0,0 +1,9 @@
+---
+'@kaizen/components': patch
+---
+
+Fix Menu dropdown not returning focus to trigger button when pressing Escape
+
+Previously, pressing Escape to close a Menu would lose focus because the focus restoration was targeting the positioning wrapper element instead of the actual button. This affected components like the Export button in TitleBlock.
+
+No changes required for consumers - this is a bug fix.

packages/components/src/MenuV1/Menu.spec.tsx

@@ -60,4 +60,26 @@ describe('Dropdown', () => {
       expect(onMouseDown).toBeCalled()
     })
   })
+
+  it('returns focus to the button when pressing Escape', async () => {
+    render(
+      <Menu button={<Button label="Button" />}>
+        <div>Item</div>
+      </Menu>,
+    )
+
+    const button = screen.getByRole('button', { name: 'Button' })
+    await user.click(button)
+
+    await waitFor(() => {
+      expect(screen.getByText('Item')).toBeVisible()
+    })
+
+    await user.keyboard('{Escape}')
+
+    await waitFor(() => {
+      expect(screen.queryByText('Item')).not.toBeInTheDocument()
+      expect(button).toHaveFocus()
+    })
+  })
 })

packages/components/src/MenuV1/Menu.tsx

@@ -1,10 +1,24 @@
-import React, { useState } from 'react'
+import React, { useState, type Ref } from 'react'
 import type { ButtonProps } from '~components/ButtonV1'
+import type { ButtonRef } from '~components/ButtonV1/GenericButton'
 import { StatelessMenu, type StatelessMenuProps } from './subcomponents/StatelessMenu'
 
+const mergeRefs =
+  <T,>(...refs: (Ref<T> | undefined)[]) =>
+  (value: T | null): void => {
+    refs.forEach((ref) => {
+      if (typeof ref === 'function') {
+        ref(value)
+      } else if (ref && typeof ref === 'object') {
+        ;(ref as React.MutableRefObject<T | null>).current = value
+      }
+    })
+  }
+
 type ButtonPropsWithOptionalAria = ButtonProps & {
   'aria-haspopup'?: boolean
   'aria-expanded'?: boolean
+  'ref'?: Ref<ButtonRef>
 }
 
 export type MenuProps = Omit<
@@ -61,6 +75,10 @@ export const Menu = ({ button, menuVisible = false, ...rest }: MenuProps): JSX.E
             props.onMouseDown(e)
             button.props.onMouseDown?.(e)
           },
+          ref: mergeRefs<ButtonRef>(
+            props.ref as Ref<ButtonRef>,
+            button.props.ref as Ref<ButtonRef> | undefined,
+          ),
         })
       }
     />

packages/components/src/MenuV1/subcomponents/MenuDropdown/MenuDropdown.tsx

@@ -18,11 +18,13 @@ export type MenuDropdownProps = {
   autoHide?: 'on' | 'outside-click-only' | 'off'
   children: React.ReactNode
   referenceElement: HTMLElement | null
+  buttonRef?: React.RefObject<{ focus: () => void } | null>
 }
 
 export const MenuDropdown = ({
   children,
   referenceElement,
+  buttonRef,
   id,
   hideMenuDropdown,
   autoHide = 'on',
@@ -104,7 +106,11 @@ export const MenuDropdown = ({
       shards={referenceElement ? [referenceElement] : undefined}
       onEscapeKey={hideMenuDropdown}
       returnFocus={() => {
-        referenceElement?.focus()
+        if (buttonRef?.current) {
+          buttonRef.current.focus()
+        } else {
+          referenceElement?.focus()
+        }
         return false
       }}
     >

packages/components/src/MenuV1/subcomponents/StatelessMenu/StatelessMenu.tsx

@@ -43,6 +43,7 @@ export type StatelessMenuProps = {
     'onClick': (e: any) => void
     'onMouseDown': (e: any) => void
     'aria-expanded': boolean
+    'ref': React.RefObject<{ focus: () => void } | null>
   }) => React.ReactElement
   'onClick'?: (event: SyntheticEvent) => void
 }
@@ -62,6 +63,7 @@ export const StatelessMenu = ({
   onClick,
 }: StatelessMenuProps): JSX.Element => {
   const [referenceElement, setReferenceElement] = useState<HTMLSpanElement | null>(null)
+  const buttonRef = useRef<{ focus: () => void } | null>(null)
   const portalSelectorElementRef = useRef<Element | null>(null)
 
   const menuButton = renderButton({
@@ -72,6 +74,7 @@ export const StatelessMenu = ({
     },
     'onMouseDown': (e: React.MouseEvent<Element, MouseEvent>) => e.preventDefault(),
     'aria-expanded': isMenuVisible,
+    'ref': buttonRef,
   })
 
   useEffect(() => {
@@ -90,6 +93,7 @@ export const StatelessMenu = ({
   const menu = isMenuVisible ? (
     <MenuDropdown
       referenceElement={referenceElement}
+      buttonRef={buttonRef}
       align={align}
       hideMenuDropdown={hideMenuDropdown}
       width={dropdownWidth}