[popups] Fix programmatic focus return

Author: atomiks

packages/react/src/dialog/root/DialogRoot.detached-triggers.test.tsx

@@ -183,7 +183,7 @@ describe('<Dialog.Root />', () => {
 
         return (
           <div>
-            <Dialog.Root open={open} triggerId={triggerId}>
+            <Dialog.Root open={open} triggerId={triggerId} onOpenChange={setOpen}>
               {({ payload }: NumberPayload) => (
                 <React.Fragment>
                   <Dialog.Trigger id="trigger-1" payload={1}>
@@ -196,6 +196,7 @@ describe('<Dialog.Root />', () => {
                   <Dialog.Portal>
                     <Dialog.Popup>
                       <span data-testid="content">{payload}</span>
+                      <Dialog.Close>Close</Dialog.Close>
                     </Dialog.Popup>
                   </Dialog.Portal>
                 </React.Fragment>
@@ -223,6 +224,13 @@ describe('<Dialog.Root />', () => {
       await waitFor(() => {
         expect(screen.getByTestId('content').textContent).toBe('2');
       });
+
+      await user.click(screen.getByRole('button', { name: 'Close' }));
+
+      await waitFor(() => {
+        expect(screen.queryByTestId('content')).toBe(null);
+      });
+      expect(openButton).toHaveFocus();
     });
 
     it('keeps the payload reactive', async () => {

packages/react/src/floating-ui-react/components/FloatingFocusManager.tsx

@@ -731,6 +731,7 @@ export function FloatingFocusManager(props: FloatingFocusManagerProps): React.JS
 
     const doc = ownerDocument(floatingFocusElement);
     const previouslyFocusedElement = activeElement(doc);
+    const openedProgrammatically = openInteractionTypeRef.current == null;
 
     addPreviouslyFocusedElement(previouslyFocusedElement);
 
@@ -793,15 +794,22 @@ export function FloatingFocusManager(props: FloatingFocusManagerProps): React.JS
         resolvedReturnFocusValue = true;
       }
 
-      if (typeof resolvedReturnFocusValue === 'boolean') {
-        if (domReference?.isConnected) {
-          return domReference;
-        }
-
-        return getPreviouslyFocusedElement() || null;
+      const connectedReference = domReference?.isConnected ? domReference : null;
+      let previouslyFocusedReturnElement = previouslyFocusedElement;
+      if (
+        !previouslyFocusedReturnElement?.isConnected ||
+        getNodeName(previouslyFocusedReturnElement) === 'body'
+      ) {
+        previouslyFocusedReturnElement = getPreviouslyFocusedElement() || null;
       }
 
-      const fallback = domReference?.isConnected ? domReference : getPreviouslyFocusedElement();
+      const fallback = openedProgrammatically
+        ? previouslyFocusedReturnElement || connectedReference
+        : connectedReference || previouslyFocusedReturnElement;
+
+      if (typeof resolvedReturnFocusValue === 'boolean') {
+        return fallback;
+      }
 
       return resolveRef(resolvedReturnFocusValue) || fallback || null;
     }
@@ -850,6 +858,7 @@ export function FloatingFocusManager(props: FloatingFocusManagerProps): React.JS
     floating,
     floatingFocusElement,
     returnFocusRef,
+    openInteractionTypeRef,
     events,
     tree,
     domReference,