[codex] fix useLockFocus retry when element is filled later

src/Dom/focus.ts

@@ -1,4 +1,5 @@
-import { useEffect } from 'react';
+import { useEffect, useRef, useState } from 'react';
+import type { DependencyList } from 'react';
 import isVisible from './isVisible';
 import useId from '../hooks/useId';
 
@@ -211,6 +212,41 @@ export function lockFocus(element: HTMLElement, id: string): VoidFunction {
   };
 }
 
+/**
+ * Retry an effect until it reports ready.
+ * When `ready` is `false`, it will schedule one more effect cycle and call `func` again
+ * with the next `retryTimes`.
+ */
+type RetryEffectResult = readonly [
+  clearFunc: VoidFunction | undefined,
+  ready: boolean,
+];
+
+function useRetryEffect(
+  func: (retryTimes: number) => RetryEffectResult,
+  deps: DependencyList,
+): void {
+  /* eslint-disable react-hooks/exhaustive-deps */
+  const retryTimesRef = useRef(0);
+  const [retryMark, setRetryMark] = useState(0);
+
+  useEffect(() => {
+    retryTimesRef.current = 0;
+  }, deps);
+
+  useEffect(() => {
+    const [clearFn, ready] = func(retryTimesRef.current);
+
+    if (!ready) {
+      retryTimesRef.current += 1;
+      setRetryMark(count => count + 1);
+    }
+
+    return clearFn;
+  }, [...deps, retryMark]);
+  /* eslint-enable react-hooks/exhaustive-deps */
+}
+
 /**
  * Lock focus within an element.
  * When locked, focus will be restricted to focusable elements within the specified element.
@@ -222,15 +258,24 @@ export function useLockFocus(
   getElement: () => HTMLElement | null,
 ): [ignoreElement: (ele: HTMLElement) => void] {
   const id = useId();
+  const getElementRef = useRef(getElement);
 
-  useEffect(() => {
-    if (lock) {
-      const element = getElement();
-      if (element) {
-        return lockFocus(element, id);
-      }
+  getElementRef.current = getElement;
+
+  const lockEffect = (retryTimes: number): RetryEffectResult => {
+    if (!lock) {
+      return [undefined, true];
     }
-  }, [lock, id]);
+
+    const element = getElementRef.current();
+    if (element) {
+      return [lockFocus(element, id), true];
+    }
+
+    return [undefined, retryTimes >= 1];
+  };
+
+  useRetryEffect(lockEffect, [id, lock]);
 
   const ignoreElement = (ele: HTMLElement) => {
     if (ele) {

tests/focus.test.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable class-methods-use-this */
-import React, { useRef } from 'react';
-import { render } from '@testing-library/react';
+import React, { useEffect, useRef, useState } from 'react';
+import { render, waitFor } from '@testing-library/react';
 import { spyElementPrototype } from '../src/test/domHook';
 import { getFocusNodeList, triggerFocus, useLockFocus } from '../src/Dom/focus';
 
@@ -95,6 +95,39 @@ describe('focus', () => {
       outerButton.focus();
       expect(document.activeElement).toBe(input1);
     });
+
+    it('should retry lock once when element is filled after lock starts', async () => {
+      const DelayedElementComponent: React.FC = () => {
+        const elementRef = useRef<HTMLDivElement>(null);
+        const [element, setElement] = useState<HTMLDivElement | null>(null);
+
+        useLockFocus(true, () => element);
+
+        useEffect(() => {
+          setElement(elementRef.current);
+        }, []);
+
+        return (
+          <>
+            <button data-testid="outer-button">Outer</button>
+            <div ref={elementRef} data-testid="focus-container" tabIndex={0}>
+              <input key="input1" data-testid="input1" />
+              <button key="button1" data-testid="button1">
+                Button
+              </button>
+            </div>
+          </>
+        );
+      };
+
+      const { getByTestId } = render(<DelayedElementComponent />);
+
+      const focusContainer = getByTestId('focus-container');
+
+      await waitFor(() => {
+        expect(document.activeElement).toBe(focusContainer);
+      });
+    });
   });
 
   it('ignoreElement should allow focus on ignored elements', () => {