fix(headless): return focus to the ancestor reference when the dialog trigger is gone

.changeset/headless-return-focus-ancestor.md

@@ -0,0 +1,2 @@
+---
+---

packages/headless/src/hooks/use-return-focus.ts

@@ -1,7 +1,8 @@
 'use client';
 
-import type { FloatingContext, OpenChangeReason } from '@floating-ui/react';
-import { useEffect, useRef } from 'react';
+import type { FloatingContext, FloatingTreeType, OpenChangeReason } from '@floating-ui/react';
+import { useFloatingParentNodeId, useFloatingTree } from '@floating-ui/react';
+import { useEffect, useMemo, useRef } from 'react';
 
 import { isKeyboardEvent } from '../utils/interaction-modality';
 
@@ -14,6 +15,12 @@ import { isKeyboardEvent } from '../utils/interaction-modality';
  * the user never asked for. A pointer dismiss therefore resolves to `null`, which leaves focus
  * where the pointer left it, the same choice Base UI makes from its close interaction type.
  *
+ * A floating element opened from inside another one — a dialog from a menu item — may have no
+ * trigger of its own, or one that is gone by the time it closes: the item unmounted with the menu.
+ * Focus then goes to the nearest ancestor in the floating tree whose reference is still on the
+ * page, which for a menu is its trigger. Resolved lazily, at restore time, since that is when it is
+ * known whether the trigger survived.
+ *
  * Pass the result to `FloatingFocusManager`'s `returnFocus`. On `null` it falls back to the
  * hidden guard element it keeps next to the trigger, so the tab position survives; verify that
  * still holds when upgrading `@floating-ui/react`.
@@ -22,12 +29,16 @@ export function useReturnFocus(
   context: Pick<FloatingContext, 'open' | 'events' | 'elements'>,
 ): React.MutableRefObject<HTMLElement | null> {
   const { open, events, elements } = context;
-  const returnFocusRef = useRef<HTMLElement | null>(null);
+  const tree = useFloatingTree();
+  const parentId = useFloatingParentNodeId();
+  const triggerRef = useRef<HTMLElement | null>(null);
+  const dismissedByPointerRef = useRef(false);
   const trigger = elements.domReference;
 
   useEffect(() => {
     if (open) {
-      returnFocusRef.current = trigger instanceof HTMLElement ? trigger : null;
+      triggerRef.current = trigger instanceof HTMLElement ? trigger : null;
+      dismissedByPointerRef.current = false;
     }
   }, [open, trigger]);
 
@@ -38,13 +49,46 @@ export function useReturnFocus(
     // closes carry no event at all.
     function onOpenChange({ open, event, reason }: { open: boolean; event?: Event; reason?: OpenChangeReason }) {
       if (!open && event && reason && !isKeyboardEvent(event)) {
-        returnFocusRef.current = null;
+        dismissedByPointerRef.current = true;
       }
     }
 
     events.on('openchange', onOpenChange);
     return () => events.off('openchange', onOpenChange);
   }, [events]);
 
-  return returnFocusRef;
+  return useMemo(
+    () => ({
+      get current() {
+        if (dismissedByPointerRef.current) {
+          return null;
+        }
+        const own = triggerRef.current;
+        if (own?.isConnected) {
+          return own;
+        }
+        return ancestorReference(tree, parentId);
+      },
+      set current(element: HTMLElement | null) {
+        triggerRef.current = element;
+      },
+    }),
+    [tree, parentId],
+  );
+}
+
+function ancestorReference(tree: FloatingTreeType | null, parentId: string | null): HTMLElement | null {
+  let id = parentId;
+  while (tree && id != null) {
+    const node = tree.nodesRef.current.find(candidate => candidate.id === id);
+    if (!node) {
+      return null;
+    }
+    const reference = node.context?.elements.domReference;
+    if (reference instanceof HTMLElement && reference.isConnected) {
+      return reference;
+    }
+    id = node.parentId ?? null;
+  }
+  return null;
 }

packages/headless/src/primitives/dialog/dialog.test.tsx

@@ -4,6 +4,7 @@ import React from 'react';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import { axe } from '../../test-utils/axe';
+import { Menu } from '../menu';
 import { Popover } from '../popover';
 import { useDialogContext } from './dialog-context';
 import { Dialog } from './index';
@@ -756,6 +757,63 @@ describe('Dialog', () => {
       expect(document.activeElement).toBe(screen.getByRole('button', { name: 'Open dialog' }));
     });
 
+    describe('opened from a menu item', () => {
+      function MenuFixture() {
+        const [open, setOpen] = React.useState(false);
+        return (
+          <Menu.Root>
+            <Menu.Trigger>Actions</Menu.Trigger>
+            <Menu.Positioner>
+              <Menu.Popup>
+                <Menu.Item
+                  label='Remove'
+                  onClick={() => setOpen(true)}
+                >
+                  Remove
+                </Menu.Item>
+              </Menu.Popup>
+            </Menu.Positioner>
+            <Dialog.Root
+              open={open}
+              onOpenChange={setOpen}
+            >
+              <Dialog.Popup>
+                <Dialog.Title>Remove?</Dialog.Title>
+                <Dialog.Close>Cancel</Dialog.Close>
+              </Dialog.Popup>
+            </Dialog.Root>
+          </Menu.Root>
+        );
+      }
+
+      it('returns focus to the menu trigger on Escape', async () => {
+        const user = userEvent.setup();
+        render(<MenuFixture />);
+
+        const trigger = screen.getByRole('button', { name: 'Actions' });
+        trigger.focus();
+        await user.keyboard('{Enter}');
+        await user.keyboard('{Enter}');
+        expect(screen.getByRole('dialog')).toBeInTheDocument();
+
+        await user.keyboard('{Escape}');
+
+        expect(document.activeElement).toBe(trigger);
+      });
+
+      it('returns focus to the menu trigger on Close press', async () => {
+        const user = userEvent.setup();
+        render(<MenuFixture />);
+
+        const trigger = screen.getByRole('button', { name: 'Actions' });
+        await user.click(trigger);
+        await user.click(screen.getByRole('menuitem', { name: 'Remove' }));
+        await user.click(screen.getByRole('button', { name: 'Cancel' }));
+
+        expect(document.activeElement).toBe(trigger);
+      });
+    });
+
     it('resolves the function form with an empty type on programmatic close', () => {
       const handle = Dialog.createHandle();
       const finalFocus = vi.fn(() => undefined);