fix(app): preserve IME composition in web overlays (#2811)

* fix(app): preserve IME composition in web overlays

* fix(app): retire detached web overlay entries

---------

Co-authored-by: northsea4 <your-email@example.com>
Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>

packages/app/src/lib/overlay-root.browser.test.tsx

@@ -0,0 +1,70 @@
+import React, { type RefCallback } from "react";
+import { flushSync } from "react-dom";
+import { createRoot, type Root } from "react-dom/client";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { useWebOverlayRegistration } from "./overlay-root";
+
+function nextFrame(): Promise<void> {
+  return new Promise((resolve) => requestAnimationFrame(() => resolve()));
+}
+
+function OverlayHarness({ active, showScope }: { active: boolean; showScope: boolean }) {
+  const setScope = useWebOverlayRegistration({ active, layer: 20, onKeyDown: () => false });
+  return showScope ? (
+    <div data-testid="scope" ref={setScope as RefCallback<HTMLDivElement>} tabIndex={-1}>
+      <input data-testid="overlay-input" />
+    </div>
+  ) : null;
+}
+
+describe("useWebOverlayRegistration in the browser", () => {
+  let root: Root;
+  let container: HTMLDivElement;
+  let opener: HTMLButtonElement;
+
+  beforeEach(() => {
+    container = document.createElement("div");
+    opener = document.createElement("button");
+    opener.textContent = "Open";
+    document.body.append(opener, container);
+    root = createRoot(container);
+    opener.focus();
+  });
+
+  afterEach(() => {
+    root.unmount();
+    document.body.replaceChildren();
+  });
+
+  async function renderHarness(active: boolean, showScope: boolean): Promise<void> {
+    flushSync(() => {
+      root.render(<OverlayHarness active={active} showScope={showScope} />);
+    });
+    await nextFrame();
+  }
+
+  it("skips focus restoration for active scope detach but restores focus when closing", async () => {
+    await renderHarness(true, true);
+
+    const input = document.querySelector<HTMLInputElement>('[data-testid="overlay-input"]');
+    expect(input).toBeTruthy();
+    input?.focus();
+    await nextFrame();
+    expect(document.activeElement).not.toBe(opener);
+
+    const originalFocus = opener.focus.bind(opener);
+    let openerFocusCalls = 0;
+    opener.focus = () => {
+      openerFocusCalls += 1;
+      originalFocus();
+    };
+
+    await renderHarness(true, false);
+
+    expect(openerFocusCalls).toBe(0);
+
+    await renderHarness(false, false);
+
+    expect(openerFocusCalls).toBe(1);
+  });
+});

packages/app/src/lib/overlay-root.test.tsx

@@ -0,0 +1,85 @@
+// @vitest-environment jsdom
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useWebOverlayRegistration } from "./overlay-root";
+
+describe("useWebOverlayRegistration", () => {
+  let opener: HTMLButtonElement;
+  let input: HTMLInputElement;
+  let scope: HTMLDivElement;
+
+  beforeEach(() => {
+    vi.spyOn(window, "requestAnimationFrame").mockImplementation((callback) => {
+      callback(0);
+      return 0;
+    });
+    vi.spyOn(window, "cancelAnimationFrame").mockImplementation(() => {});
+
+    opener = document.createElement("button");
+    input = document.createElement("input");
+    scope = document.createElement("div");
+    scope.tabIndex = -1;
+    scope.append(input);
+    document.body.append(opener, scope);
+    opener.focus();
+    expect(document.activeElement).toBe(opener);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    document.body.replaceChildren();
+  });
+
+  it("does not restore opener focus when an active overlay scope detaches", () => {
+    const { result, unmount } = renderHook(() =>
+      useWebOverlayRegistration({ active: true, layer: 20, onKeyDown: () => false }),
+    );
+
+    act(() => result.current(scope));
+    const openerFocus = vi.spyOn(opener, "focus");
+    input.focus();
+    expect(document.activeElement).toBe(input);
+
+    act(() => result.current(null));
+
+    expect(openerFocus).not.toHaveBeenCalled();
+    expect(document.activeElement).toBe(input);
+    unmount();
+  });
+
+  it("restores opener focus when an overlay closes", () => {
+    const { result, rerender, unmount } = renderHook(
+      ({ active }: { active: boolean }) =>
+        useWebOverlayRegistration({ active, layer: 20, onKeyDown: () => false }),
+      { initialProps: { active: true } },
+    );
+
+    act(() => result.current(scope));
+    const openerFocus = vi.spyOn(opener, "focus");
+    input.focus();
+    expect(document.activeElement).toBe(input);
+
+    act(() => rerender({ active: false }));
+
+    expect(openerFocus).toHaveBeenCalled();
+    unmount();
+  });
+
+  it("restores opener focus when an active overlay unmounts after its scope detaches", () => {
+    const { result, unmount } = renderHook(() =>
+      useWebOverlayRegistration({ active: true, layer: 20, onKeyDown: () => false }),
+    );
+
+    act(() => result.current(scope));
+    const openerFocus = vi.spyOn(opener, "focus");
+    input.focus();
+    expect(document.activeElement).toBe(input);
+
+    act(() => result.current(null));
+    expect(openerFocus).not.toHaveBeenCalled();
+
+    unmount();
+
+    expect(openerFocus).toHaveBeenCalled();
+  });
+});

packages/app/src/lib/overlay-root.ts

@@ -95,6 +95,10 @@ let webOverlayOrder = 0;
 let webOverlayListenersAttached = false;
 let webOverlayFocusCheckQueued = false;
 
+interface RemoveWebOverlayOptions {
+  restoreFocus?: boolean;
+}
+
 function getTopWebOverlay(): WebOverlayEntry | undefined {
   return webOverlayEntries.reduce<WebOverlayEntry | undefined>((top, entry) => {
     if (!top) return entry;
@@ -195,7 +199,7 @@ function detachWebOverlayListeners(): void {
   webOverlayListenersAttached = false;
 }
 
-function addWebOverlay(entry: WebOverlayEntry): () => void {
+function addWebOverlay(entry: WebOverlayEntry): (options?: RemoveWebOverlayOptions) => void {
   webOverlayEntries.push(entry);
   attachWebOverlayListeners();
 
@@ -206,12 +210,16 @@ function addWebOverlay(entry: WebOverlayEntry): () => void {
     }
   });
 
-  return () => {
+  return (options) => {
     window.cancelAnimationFrame(focusFrame);
     const index = webOverlayEntries.findIndex((candidate) => candidate.id === entry.id);
     if (index !== -1) webOverlayEntries.splice(index, 1);
     detachWebOverlayListeners();
-    if (entry.restoreFocus && document.contains(entry.restoreFocus)) {
+    if (
+      options?.restoreFocus !== false &&
+      entry.restoreFocus &&
+      document.contains(entry.restoreFocus)
+    ) {
       entry.restoreFocus.focus();
     }
   };
@@ -240,7 +248,9 @@ export function useWebOverlayRegistration({
   const layerRef = useRef(layer);
   const keyHandlerRef = useRef(onKeyDown);
   const capturedRestoreFocusRef = useRef<HTMLElement | null>(null);
-  const removeEntryRef = useRef<(() => void) | null>(null);
+  const removeEntryRef = useRef<((options?: RemoveWebOverlayOptions) => void) | null>(null);
+  const detachedRemoveEntryRef = useRef<((options?: RemoveWebOverlayOptions) => void) | null>(null);
+  const registeredRef = useRef(false);
   const activeRef = useRef(active);
   const wasActiveRef = useRef(false);
 
@@ -265,12 +275,20 @@ export function useWebOverlayRegistration({
       typeof window !== "undefined" &&
       typeof document !== "undefined";
     if (!shouldRegister) {
-      removeEntryRef.current?.();
+      const shouldRestoreFocus = !activeRef.current;
-      removeEntryRef.current = null;
+      if (registeredRef.current || shouldRestoreFocus) {
+        removeEntryRef.current?.({ restoreFocus: shouldRestoreFocus });
+      }
+      registeredRef.current = false;
+      if (shouldRestoreFocus) {
+        removeEntryRef.current = null;
+        detachedRemoveEntryRef.current = null;
+      }
       return;
     }
-    if (removeEntryRef.current) return;
+    if (registeredRef.current) return;
 
+    detachedRemoveEntryRef.current = null;
     const entry: WebOverlayEntry = {
       id: idRef.current,
       order: ++webOverlayOrder,
@@ -280,12 +298,20 @@ export function useWebOverlayRegistration({
       restoreFocus: capturedRestoreFocusRef.current,
     };
     removeEntryRef.current = addWebOverlay(entry);
+    registeredRef.current = true;
   }, []);
 
   const setScope = useCallback(
     (node: unknown) => {
       scopeRef.current =
         typeof HTMLElement !== "undefined" && node instanceof HTMLElement ? node : null;
+      if (!scopeRef.current && activeRef.current) {
+        const removeEntry = removeEntryRef.current;
+        removeEntry?.({ restoreFocus: false });
+        detachedRemoveEntryRef.current = removeEntry;
+        removeEntryRef.current = null;
+        registeredRef.current = false;
+      }
       // Host refs attach before descendant layout effects and autofocus. Register
       // here so the previous overlay cannot redirect that pending focus.
       syncRegistration();
@@ -296,8 +322,11 @@ export function useWebOverlayRegistration({
   useLayoutEffect(() => {
     syncRegistration();
     return () => {
-      removeEntryRef.current?.();
+      const removeEntry = removeEntryRef.current ?? detachedRemoveEntryRef.current;
+      removeEntry?.();
       removeEntryRef.current = null;
+      detachedRemoveEntryRef.current = null;
+      registeredRef.current = false;
     };
   }, [active, syncRegistration]);