fix: restore selection should consider the window of the container

packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js

@@ -325,7 +325,7 @@ export function getPublicInstance(instance: Instance): Instance {
 
 export function prepareForCommit(containerInfo: Container): Object | null {
   eventsEnabled = ReactBrowserEventEmitterIsEnabled();
-  selectionInformation = getSelectionInformation();
+  selectionInformation = getSelectionInformation(containerInfo);
   let activeInstance = null;
   if (enableCreateEventHandleAPI) {
     const focusedElem = selectionInformation.focusedElem;
@@ -357,7 +357,7 @@ export function afterActiveInstanceBlur(): void {
 }
 
 export function resetAfterCommit(containerInfo: Container): void {
-  restoreSelection(selectionInformation);
+  restoreSelection(selectionInformation, containerInfo);
   ReactBrowserEventEmitterSetEnabled(eventsEnabled);
   eventsEnabled = null;
   selectionInformation = null;

packages/react-dom-bindings/src/client/ReactInputSelection.js

@@ -56,9 +56,9 @@ function isSameOriginFrame(iframe) {
   }
 }
 
-function getActiveElementDeep() {
-  let win = window;
-  let element = getActiveElement();
+function getActiveElementDeep(containerInfo) {
+  let win = containerInfo?.ownerDocument?.defaultView ?? window;
+  let element = getActiveElement(win.document);
   while (element instanceof win.HTMLIFrameElement) {
     if (isSameOriginFrame(element)) {
       win = element.contentWindow;
@@ -97,8 +97,8 @@ export function hasSelectionCapabilities(elem) {
   );
 }
 
-export function getSelectionInformation() {
-  const focusedElem = getActiveElementDeep();
+export function getSelectionInformation(containerInfo) {
+  const focusedElem = getActiveElementDeep(containerInfo);
   return {
     focusedElem: focusedElem,
     selectionRange: hasSelectionCapabilities(focusedElem)
@@ -112,8 +112,8 @@ export function getSelectionInformation() {
  * restore it. This is useful when performing operations that could remove dom
  * nodes and place them back in, resulting in focus being lost.
  */
-export function restoreSelection(priorSelectionInformation) {
-  const curFocusedElem = getActiveElementDeep();
+export function restoreSelection(priorSelectionInformation, containerInfo) {
+  const curFocusedElem = getActiveElementDeep(containerInfo);
   const priorFocusedElem = priorSelectionInformation.focusedElem;
   const priorSelectionRange = priorSelectionInformation.selectionRange;
   if (curFocusedElem !== priorFocusedElem && isInDocument(priorFocusedElem)) {

packages/react-dom/src/__tests__/ReactDOMFiber-test.js

@@ -19,12 +19,23 @@ let act;
 let assertConsoleErrorDev;
 let assertLog;
 let root;
+let JSDOM
 
 describe('ReactDOMFiber', () => {
   let container;
 
   beforeEach(() => {
     jest.resetModules();
+
+    // JSDOM needs to be setup with a TextEncoder and TextDecoder when used standalone
+    // https://github.com/jsdom/jsdom/issues/2524
+    (() => {
+      const { TextEncoder, TextDecoder } = require('util');
+      global.TextEncoder = TextEncoder;
+      global.TextDecoder = TextDecoder;
+      JSDOM = require('jsdom').JSDOM;
+    })();
+
     React = require('react');
     ReactDOM = require('react-dom');
     PropTypes = require('prop-types');
@@ -1272,4 +1283,48 @@ describe('ReactDOMFiber', () => {
     });
     expect(didCallOnChange).toBe(true);
   });
+
+  it('should restore selection in the correct window', async () => {
+    // creating new JSOM instance to get a second window as window.open is not implemented
+    // https://github.com/jsdom/jsdom/blob/c53efc81e75f38a0558fbf3ed75d30b78b4c4898/lib/jsdom/browser/Window.js#L987
+    const {window: newWindow} = new JSDOM('');
+    // creating a new container since the default cleanup expects the existing container to be in the document
+    const newContainer = newWindow.document.createElement('div');
+    newWindow.document.body.appendChild(newContainer);
+    root = ReactDOMClient.createRoot(newContainer);
+
+    const Test = () => {
+      const [reverse, setReverse] = React.useState(false);
+      const [items] = React.useState(() => ['a', 'b', 'c']);
+      const onClick = () => {
+        setReverse(true);
+      };
+
+      // shuffle the items so that the react commit needs to restore focus
+      // to the correct element after commit
+      const itemsToRender = reverse ? items.reverse() : items;
+
+      return (
+        <div>
+          {itemsToRender.map(item => (
+            <button onClick={onClick} key={item} id={item}>
+              {item}
+            </button>
+          ))}
+        </div>
+      );
+    };
+
+    await act(() => {
+      root.render(<Test />);
+    });
+
+    newWindow.document.getElementById('a').focus();
+    await act(() => {
+      newWindow.document.getElementById('a').click();
+    });
+
+    expect(newWindow.document.activeElement).not.toBe(newWindow.document.body);
+    expect(newWindow.document.activeElement.innerHTML).toBe('a');
+  });
 });