From 6466591e019c8d0f9841e011b4fbabe90039de87 Mon Sep 17 00:00:00 2001
From: Jerry Jones <jones.jeremydavid@gmail.com>
Date: Wed, 18 Dec 2024 12:30:01 -0600
Subject: [PATCH] useFocusReturn: Get active element within the iframe when
 restoring focus (#68060)

Co-authored-by: Alex Lende <alex+github.com@lende.xyz>
---
 .../src/components/button-block-appender/index.js        | 9 ++-------
 packages/compose/src/hooks/use-focus-return/index.js     | 8 +++++++-
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/packages/block-editor/src/components/button-block-appender/index.js b/packages/block-editor/src/components/button-block-appender/index.js
index 53b15e2fd2cfdd..4cde8c26d75638 100644
--- a/packages/block-editor/src/components/button-block-appender/index.js
+++ b/packages/block-editor/src/components/button-block-appender/index.js
@@ -7,7 +7,7 @@ import clsx from 'clsx';
  * WordPress dependencies
  */
 import { Button } from '@wordpress/components';
-import { forwardRef, useRef } from '@wordpress/element';
+import { forwardRef } from '@wordpress/element';
 import { _x, sprintf } from '@wordpress/i18n';
 import { Icon, plus } from '@wordpress/icons';
 import deprecated from '@wordpress/deprecated';
@@ -16,15 +16,11 @@ import deprecated from '@wordpress/deprecated';
  * Internal dependencies
  */
 import Inserter from '../inserter';
-import { useMergeRefs } from '@wordpress/compose';
 
 function ButtonBlockAppender(
 	{ rootClientId, className, onFocus, tabIndex, onSelect },
 	ref
 ) {
-	const inserterButtonRef = useRef();
-
-	const mergedInserterButtonRef = useMergeRefs( [ inserterButtonRef, ref ] );
 	return (
 		<Inserter
 			position="bottom center"
@@ -34,7 +30,6 @@ function ButtonBlockAppender(
 				if ( onSelect && typeof onSelect === 'function' ) {
 					onSelect( ...args );
 				}
-				inserterButtonRef.current?.focus();
 			} }
 			renderToggle={ ( {
 				onToggle,
@@ -61,7 +56,7 @@ function ButtonBlockAppender(
 				return (
 					<Button
 						__next40pxDefaultSize
-						ref={ mergedInserterButtonRef }
+						ref={ ref }
 						onFocus={ onFocus }
 						tabIndex={ tabIndex }
 						className={ clsx(
diff --git a/packages/compose/src/hooks/use-focus-return/index.js b/packages/compose/src/hooks/use-focus-return/index.js
index 2cd93b279cd318..36dc7560669652 100644
--- a/packages/compose/src/hooks/use-focus-return/index.js
+++ b/packages/compose/src/hooks/use-focus-return/index.js
@@ -48,7 +48,13 @@ function useFocusReturn( onFocusReturn ) {
 				return;
 			}
 
-			focusedBeforeMount.current = node.ownerDocument.activeElement;
+			const activeDocument =
+				node.ownerDocument.activeElement instanceof
+				window.HTMLIFrameElement
+					? node.ownerDocument.activeElement.contentDocument
+					: node.ownerDocument;
+
+			focusedBeforeMount.current = activeDocument?.activeElement ?? null;
 		} else if ( focusedBeforeMount.current ) {
 			const isFocused = ref.current?.contains(
 				ref.current?.ownerDocument.activeElement