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,24 +134,22 @@ 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) {
+      } else if (document.activeElement === lastElement) {
         e.preventDefault();
         firstElement?.focus();
       }
-      }
     };
 
     document.addEventListener("keydown", handleTab);
@@ -76,7 +157,7 @@ export function useCreateModalA11y(
     return () => {
       document.body.style.overflow = "";
       document.removeEventListener("keydown", handleTab);
-      previousActiveElementRef.current?.focus();
+      restoreFocus(previousActiveElementRef.current);
     };
   }, [dialogRef, isOpen]);
 }

tests/components/Create.test.tsx

@@ -1,6 +1,7 @@
-import React from "react";
+import React, { useState } from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import { screen, fireEvent, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import "@testing-library/jest-dom/vitest";
 import { renderWithProviders } from "../utils/test-utils";
 import Create from "../../app/components/modals/Create";
@@ -8,6 +9,23 @@ import TextInput from "../../app/components/controls/TextInput";
 
 type CreateProps = React.ComponentProps<typeof Create>;
 
+function CreateOpenHarness() {
+  const [open, setOpen] = useState(false);
+  return (
+    <>
+      <button type="button" onClick={() => setOpen(true)}>
+        Open create dialog
+      </button>
+      <Create
+        isOpen={open}
+        onClose={() => setOpen(false)}
+        title="Test Create Dialog"
+        description="Test description"
+      />
+    </>
+  );
+}
+
 describe("Create", () => {
   const defaultProps: CreateProps = {
     isOpen: true,
@@ -233,6 +251,17 @@ describe("Create", () => {
     expect(document.body.style.overflow).toBe("");
   });
 
+  it("restores focus to the control that opened it", async () => {
+    const user = userEvent.setup();
+    renderWithProviders(<CreateOpenHarness />);
+    const trigger = screen.getByRole("button", { name: "Open create dialog" });
+    await user.click(trigger);
+    expect(screen.getByRole("dialog")).toBeInTheDocument();
+    await user.keyboard("{Escape}");
+    expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
+    expect(trigger).toHaveFocus();
+  });
+
   it("traps focus within create dialog", async () => {
     renderWithProviders(
       <Create {...defaultProps}>

tests/components/Dialog.test.tsx

@@ -1,12 +1,84 @@
-import React from "react";
+import React, { useState } from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import { screen, fireEvent } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import "@testing-library/jest-dom/vitest";
 import { renderWithProviders } from "../utils/test-utils";
 import Dialog from "../../app/components/modals/Dialog";
 
 type Props = React.ComponentProps<typeof Dialog>;
 
+const dialogCopy = {
+  title: "Confirm action",
+  description: "This cannot be undone.",
+};
+
+function DialogOpenHarness() {
+  const [open, setOpen] = useState(false);
+  return (
+    <>
+      <button type="button" onClick={() => setOpen(true)}>
+        Open dialog
+      </button>
+      <Dialog
+        isOpen={open}
+        onClose={() => setOpen(false)}
+        title={dialogCopy.title}
+        description={dialogCopy.description}
+        footer={
+          <button type="button" onClick={() => setOpen(false)}>
+            Cancel
+          </button>
+        }
+      />
+    </>
+  );
+}
+
+function DialogFromMenuHarness() {
+  const [menuOpen, setMenuOpen] = useState(false);
+  const [dialogOpen, setDialogOpen] = useState(false);
+  const menuId = "dialog-focus-restore-menu";
+  return (
+    <>
+      <button
+        type="button"
+        aria-haspopup="menu"
+        aria-expanded={menuOpen}
+        aria-controls={menuId}
+        onClick={() => setMenuOpen((open) => !open)}
+      >
+        More options
+      </button>
+      {menuOpen ? (
+        <div role="menu" id={menuId}>
+          <button
+            type="button"
+            role="menuitem"
+            onClick={() => {
+              setDialogOpen(true);
+              setMenuOpen(false);
+            }}
+          >
+            Remove
+          </button>
+        </div>
+      ) : null}
+      <Dialog
+        isOpen={dialogOpen}
+        onClose={() => setDialogOpen(false)}
+        title={dialogCopy.title}
+        description={dialogCopy.description}
+        footer={
+          <button type="button" onClick={() => setDialogOpen(false)}>
+            Cancel
+          </button>
+        }
+      />
+    </>
+  );
+}
+
 describe("Dialog", () => {
   const defaultProps: Props = {
     isOpen: true,
@@ -57,4 +129,41 @@ describe("Dialog", () => {
     renderWithProviders(<Dialog {...defaultProps} />);
     expect(document.body.style.overflow).toBe("hidden");
   });
+
+  it("restores focus to the control that opened it", async () => {
+    const user = userEvent.setup();
+    renderWithProviders(<DialogOpenHarness />);
+    const trigger = screen.getByRole("button", { name: "Open dialog" });
+    await user.click(trigger);
+    expect(screen.getByRole("dialog")).toBeInTheDocument();
+    await user.keyboard("{Escape}");
+    expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
+    expect(trigger).toHaveFocus();
+  });
+
+  it("restores focus to the opener when the opening click did not focus it", () => {
+    renderWithProviders(<DialogOpenHarness />);
+    const trigger = screen.getByRole("button", { name: "Open dialog" });
+    fireEvent.pointerDown(trigger);
+    fireEvent.click(trigger);
+    expect(screen.getByRole("dialog")).toBeInTheDocument();
+    fireEvent.keyDown(document, { key: "Escape" });
+    expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
+    expect(trigger).toHaveFocus();
+  });
+
+  it("restores focus to the menu trigger when a menu item opened the dialog", async () => {
+    const user = userEvent.setup();
+    renderWithProviders(<DialogFromMenuHarness />);
+    const menuTrigger = screen.getByRole("button", { name: "More options" });
+    await user.click(menuTrigger);
+    await user.click(screen.getByRole("menuitem", { name: "Remove" }));
+    expect(screen.getByRole("dialog")).toBeInTheDocument();
+    expect(
+      screen.queryByRole("menuitem", { name: "Remove" }),
+    ).not.toBeInTheDocument();
+    await user.keyboard("{Escape}");
+    expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
+    expect(menuTrigger).toHaveFocus();
+  });
 });