Return focus to the control that opened a dialog instead of leaving it on the document body.

Co-authored-by: Cursor <cursoragent@cursor.com>

app/components/modals/Create/useCreateModalA11y.ts

@@ -1,10 +1,93 @@
 "use client";
 
 import type { RefObject } from "react";
-import { useEffect, useRef } from "react";
+import { useEffect, useLayoutEffect, useRef } from "react";
+
+const FOCUSABLE_SELECTOR =
+  'button, [href], input, select, textarea, [tabindex]:not([tabindex="-1"])';
+
+let lastInteractedFocusable: HTMLElement | null = null;
+let interactionTrackingBound = false;
+
+function closestFocusable(target: EventTarget | null): HTMLElement | null {
+  if (!(target instanceof Element)) return null;
+  const match = target.closest(FOCUSABLE_SELECTOR);
+  return match instanceof HTMLElement ? match : null;
+}
+
+/** Menu items unmount on select; restore to the control that opened the menu. */
+function menuTriggerFor(item: HTMLElement): HTMLElement | null {
+  if (item.getAttribute("role") !== "menuitem") return null;
+  const menu = item.closest("[role='menu']");
+  const menuId = menu?.getAttribute("id");
+  if (menuId) {
+    const trigger = document.querySelector(
+      `[aria-controls="${CSS.escape(menuId)}"]`,
+    );
+    if (trigger instanceof HTMLElement) return trigger;
+  }
+  const expanded = document.querySelector(
+    '[aria-haspopup="menu"][aria-expanded="true"]',
+  );
+  return expanded instanceof HTMLElement ? expanded : null;
+}
+
+function stableFocusableFrom(target: EventTarget | null): HTMLElement | null {
+  const focusable = closestFocusable(target);
+  if (!focusable) return null;
+  return menuTriggerFor(focusable) ?? focusable;
+}
+
+function isRestorable(
+  node: HTMLElement | null,
+  dialog: HTMLElement | null,
+): node is HTMLElement {
+  if (!node?.isConnected) return false;
+  if (node === document.body || node === document.documentElement) return false;
+  if (dialog?.contains(node)) return false;
+  return true;
+}
+
+function retainLastInteracted(): void {
+  if (interactionTrackingBound || typeof document === "undefined") return;
+  interactionTrackingBound = true;
+  const save = (event: Event) => {
+    const el = stableFocusableFrom(event.target);
+    if (!el) return;
+    lastInteractedFocusable = el;
+  };
+  document.addEventListener("pointerdown", save, true);
+  document.addEventListener("focusin", save);
+}
+
+retainLastInteracted();
+
+function snapshotTrigger(dialog: HTMLElement | null): HTMLElement | null {
+  if (lastInteractedFocusable && !lastInteractedFocusable.isConnected) {
+    lastInteractedFocusable = null;
+  }
+  const active =
+    document.activeElement instanceof HTMLElement
+      ? document.activeElement
+      : null;
+  const fromActive = isRestorable(active, dialog)
+    ? (menuTriggerFor(active) ?? active)
+    : null;
+  if (fromActive && isRestorable(fromActive, dialog)) return fromActive;
+  if (isRestorable(lastInteractedFocusable, dialog)) {
+    return lastInteractedFocusable;
+  }
+  return null;
+}
+
+function restoreFocus(node: HTMLElement | null): void {
+  if (!node?.isConnected) return;
+  node.focus();
+}
 
 /**
- * Escape-to-close, body scroll lock, focus move-in and tab trap for Create-shell modals.
+ * Escape-to-close, body scroll lock, focus move-in, tab trap, and restore
+ * focus to the control that opened a Create-shell modal.
  */
 export function useCreateModalA11y(
   isOpen: boolean,
@@ -28,17 +111,17 @@ export function useCreateModalA11y(
     };
   }, [isOpen, onClose]);
 
-  useEffect(() => {
+  useLayoutEffect(() => {
     if (!isOpen) return;
 
-    previousActiveElementRef.current = document.activeElement as HTMLElement;
+    previousActiveElementRef.current = snapshotTrigger(dialogRef.current);
     document.body.style.overflow = "hidden";
 
     if (dialogRef.current) {
       const focusableElements = dialogRef.current.querySelectorAll(
-        'button, [href], input, select, textarea, [tabindex]:not([tabindex="-1"])',
+        FOCUSABLE_SELECTOR,
       );
-      const firstElement = focusableElements[0] as HTMLElement;
+      const firstElement = focusableElements[0] as HTMLElement | undefined;
       if (firstElement) {
         firstElement.focus();
       } else {
@@ -51,23 +134,21 @@ export function useCreateModalA11y(
       if (e.key !== "Tab" || !dialogRef.current) return;
 
       const focusableElements = dialogRef.current.querySelectorAll(
-        'button, [href], input, select, textarea, [tabindex]:not([tabindex="-1"])',
+        FOCUSABLE_SELECTOR,
       );
-      const firstElement = focusableElements[0] as HTMLElement;
+      const firstElement = focusableElements[0] as HTMLElement | undefined;
       const lastElement = focusableElements[
         focusableElements.length - 1
-      ] as HTMLElement;
+      ] as HTMLElement | undefined;
 
       if (e.shiftKey) {
         if (document.activeElement === firstElement) {
           e.preventDefault();
           lastElement?.focus();
         }
-      } else {
-        if (document.activeElement === lastElement) {
-          e.preventDefault();
-          firstElement?.focus();
-        }
+      } else if (document.activeElement === lastElement) {
+        e.preventDefault();
+        firstElement?.focus();
       }
     };
 
@@ -76,7 +157,7 @@ export function useCreateModalA11y(
     return () => {
       document.body.style.overflow = "";
       document.removeEventListener("keydown", handleTab);
-      previousActiveElementRef.current?.focus();
+      restoreFocus(previousActiveElementRef.current);
     };
   }, [dialogRef, isOpen]);
 }