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

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

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();
+  });
 });