Keep focus where the user moved it when a menu finishes closing

src/shared/design-system/DESIGN.md

@@ -292,6 +292,10 @@ close button and actions. Pending operations set preventClose so Escape and the
 close button agree. It retains the app's explicit dismissal behavior: outside
 clicks do not discard a form. Provide initialFocus for search dialogs and
 finalFocus when a flow has an external trigger or opens a second dialog.
+An explicit finalFocus applies only while focus is still in the closing popup
+or on the page body. If the user already moved focus elsewhere, closing leaves
+it there, as the default `true` does (`ui/finalFocus.ts`, shared by Menu,
+Popover and Dialog).
 Editors can supply `headerActions` beside Close and `leadingActions` before the
 trailing footer actions. `onEscape` may return true to consume Escape for an
 inline layer (such as an inspector) before dismissing the dialog. Nested modal

src/shared/design-system/ui/Dialog.tsx

@@ -2,6 +2,7 @@ import { Dialog as BaseDialog } from "@base-ui/react/dialog";
 import { XIcon } from "../icons";
 import { useState, type ComponentProps, type ReactNode } from "react";
 import { IconButton } from "./IconButton";
+import { useFinalFocusUnlessMoved } from "./finalFocus";
 
 type PopupProps = ComponentProps<typeof BaseDialog.Popup>;
 export type DialogProps = {
@@ -56,6 +57,7 @@ export function Dialog({
   finalFocus,
 }: DialogProps) {
   const [instantClose, setInstantClose] = useState(false);
+  const focus = useFinalFocusUnlessMoved(finalFocus, undefined);
   const transition =
     motion === "none" || (!open && instantClose) ? "none" : "default";
   return (
@@ -93,7 +95,8 @@ export function Dialog({
           data-motion={transition}
           aria-modal="true"
           initialFocus={initialFocus}
-          finalFocus={finalFocus}
+          ref={focus.ref}
+          finalFocus={focus.finalFocus}
         >
           <header className="buzz-dialog-header">
             <div className="buzz-dialog-heading">

src/shared/design-system/ui/Menu.tsx

@@ -1,6 +1,7 @@
 import { ContextMenu as BaseContextMenu } from "@base-ui/react/context-menu";
 import { Menu as BaseMenu } from "@base-ui/react/menu";
 import { CheckIcon, CaretRightIcon } from "../icons/index";
+import { useFinalFocusUnlessMoved } from "./finalFocus";
 import type { ComponentProps, ReactNode } from "react";
 
 export const MenuRoot = BaseMenu.Root;
@@ -40,8 +41,11 @@ function PositionedPopup({
   side = "bottom",
   sideOffset = 4,
   sticky,
+  finalFocus,
+  ref,
   ...props
 }: PopupProps) {
+  const focus = useFinalFocusUnlessMoved(finalFocus, ref);
   return (
     <BaseMenu.Portal>
       <BaseMenu.Positioner
@@ -57,6 +61,8 @@ function PositionedPopup({
       >
         <BaseMenu.Popup
           {...props}
+          ref={focus.ref}
+          finalFocus={focus.finalFocus}
           data-buzz-ui=""
           data-size={size}
           className="buzz-menu-popup text-body-sm"

src/shared/design-system/ui/Popover.tsx

@@ -1,5 +1,6 @@
 import { Popover as BasePopover } from "@base-ui/react/popover";
 import type { ComponentProps } from "react";
+import { useFinalFocusUnlessMoved } from "./finalFocus";
 
 export const PopoverRoot = BasePopover.Root;
 export const PopoverTrigger = BasePopover.Trigger;
@@ -42,8 +43,11 @@ export function PopoverPopup({
   size = "default",
   padding = "content",
   colorMode,
+  finalFocus,
+  ref,
   ...props
 }: PopoverPopupProps) {
+  const focus = useFinalFocusUnlessMoved(finalFocus, ref);
   return (
     <BasePopover.Portal>
       <BasePopover.Positioner
@@ -59,6 +63,8 @@ export function PopoverPopup({
       >
         <BasePopover.Popup
           {...props}
+          ref={focus.ref}
+          finalFocus={focus.finalFocus}
           data-buzz-ui=""
           data-size={size}
           data-padding={padding}

src/shared/design-system/ui/finalFocus.ts

@@ -0,0 +1,76 @@
+import type { Menu } from "@base-ui/react/menu";
+import {
+  useCallback,
+  useRef,
+  type ComponentProps,
+  type Ref,
+  type RefCallback,
+} from "react";
+
+/** Menu, Popover and Dialog popups share this Base UI prop. */
+type FinalFocus = ComponentProps<typeof Menu.Popup>["finalFocus"];
+type CloseType = Parameters<
+  Extract<FinalFocus, (...args: never) => unknown>
+>[0];
+
+/**
+ * Focus inside another closing popup (a submenu, or the popup this one is
+ * nested in) is part of the closing tree, not a user move. Collapsed
+ * disclosures also carry `data-closed`, so match popup roles only.
+ */
+const closingPopup =
+  ':is([role="menu"], [role="dialog"], [role="alertdialog"], [role="listbox"])[data-closed]';
+
+/**
+ * Keep an explicit `finalFocus` from pulling focus back after the user moved it.
+ *
+ * Base UI returns focus when a popup finishes closing. For the default `true`
+ * it first checks whether focus already moved somewhere else, but an explicit
+ * target skips that check. A user who pressed Escape and at once clicked
+ * another menu button then lost that new menu: the old popup finished its
+ * exit, focused its own trigger, and the new menu closed on focus loss.
+ *
+ * The popup element is kept after unmount so the check still works when Base
+ * UI resolves focus after React clears the ref.
+ */
+export function useFinalFocusUnlessMoved(
+  finalFocus: FinalFocus,
+  ref: Ref<HTMLDivElement> | undefined,
+): { ref: RefCallback<HTMLDivElement>; finalFocus: FinalFocus } {
+  const popup = useRef<HTMLDivElement | null>(null);
+  const setPopup = useCallback<RefCallback<HTMLDivElement>>(
+    (element: HTMLDivElement | null) => {
+      if (element) popup.current = element;
+      if (typeof ref === "function") return ref(element);
+      if (ref) ref.current = element;
+    },
+    [ref],
+  );
+  const resolve = useCallback(
+    (closeType: CloseType) => {
+      const target =
+        typeof finalFocus === "function"
+          ? finalFocus(closeType)
+          : typeof finalFocus === "object"
+            ? finalFocus.current
+            : finalFocus;
+      const doc = popup.current?.ownerDocument ?? document;
+      const active = doc.activeElement;
+      const moved =
+        active instanceof HTMLElement &&
+        active !== doc.body &&
+        active !== target &&
+        !popup.current?.contains(active) &&
+        !active.closest(closingPopup);
+      return moved ? false : target;
+    },
+    [finalFocus],
+  );
+  return {
+    ref: setPopup,
+    finalFocus:
+      finalFocus === undefined || typeof finalFocus === "boolean"
+        ? finalFocus
+        : resolve,
+  };
+}

tests/browser/menu-dismiss.spec.mjs

@@ -0,0 +1,114 @@
+import { test, expect } from "./fixture.mjs";
+import { open } from "./timeline.mjs";
+
+// Browser-only boundary: Base UI restores focus when a popup's exit
+// transition ends. The test holds that transition, so the user's next focus
+// move lands while the old menu is still closing.
+async function holdMenuExit(page) {
+  await page.evaluate(() => {
+    document.addEventListener(
+      "transitionrun",
+      (event) => {
+        if (
+          event.target instanceof HTMLElement &&
+          event.target.matches(".buzz-menu-popup[data-ending-style]")
+        )
+          for (const animation of event.target.getAnimations())
+            animation.pause();
+      },
+      true,
+    );
+  });
+  return {
+    // Closing has started and cannot finish until release().
+    held: (popup) =>
+      expect
+        .poll(() =>
+          popup.evaluate((element) =>
+            element.getAnimations().some((a) => a.playState === "paused"),
+          ),
+        )
+        .toBe(true),
+    release: () =>
+      page.evaluate(() => {
+        for (const animation of document.getAnimations())
+          if (animation.playState === "paused") animation.finish();
+      }),
+  };
+}
+
+// The entrance has started and settled. Animations alone are not enough: an
+// Escape before the entrance transition starts closes without any transition.
+async function opened(popup) {
+  await expect(popup).toHaveAttribute("data-open", "");
+  await expect(popup).not.toHaveAttribute("data-starting-style");
+  await expect
+    .poll(() => popup.evaluate((element) => element.getAnimations().length))
+    .toBe(0);
+}
+
+test("closing a menu keeps focus where the user moved it", async ({
+  page,
+  app,
+}) => {
+  await open(page, app);
+  const exit = await holdMenuExit(page);
+  const row = page
+    .getByRole("navigation", {
+      name: "Subscribed channels",
+      includeHidden: true,
+    })
+    .locator('[data-channel-id="beta"]');
+  const rowMenu = page.getByRole("menu", { name: "Actions for Beta" });
+  const composer = page.getByRole("textbox", {
+    name: "Message #Alpha",
+    exact: true,
+  });
+
+  try {
+    await row.click({ button: "right" });
+    await opened(rowMenu);
+    await page.keyboard.press("Escape");
+    await exit.held(rowMenu);
+    await composer.click();
+    await expect(composer).toBeFocused();
+    await expect(rowMenu).toHaveCount(1);
+    await exit.release();
+    // The row menu has finished closing. Its explicit final focus (the row)
+    // must not pull focus back out of the composer.
+    await expect(rowMenu).toHaveCount(0);
+    await expect(composer).toBeFocused();
+
+    // A collapsed disclosure (Base UI marks it data-closed, like a closing
+    // popup) is ordinary page content. Focus moved into it also stays.
+    await page.evaluate(() => {
+      const collapsed = document.createElement("div");
+      collapsed.setAttribute("data-closed", "");
+      collapsed.innerHTML = '<button type="button">Collapsed control</button>';
+      document.getElementById("main-content").append(collapsed);
+    });
+    const collapsedControl = page.getByRole("button", {
+      name: "Collapsed control",
+    });
+    await row.click({ button: "right" });
+    await opened(rowMenu);
+    await page.keyboard.press("Escape");
+    await exit.held(rowMenu);
+    await collapsedControl.focus();
+    await expect(rowMenu).toHaveCount(1);
+    await exit.release();
+    await expect(rowMenu).toHaveCount(0);
+    await expect(collapsedControl).toBeFocused();
+
+    // With nothing else focused, Escape still returns focus to the row.
+    await row.click({ button: "right" });
+    await opened(rowMenu);
+    await page.keyboard.press("Escape");
+    await exit.held(rowMenu);
+    await exit.release();
+    await expect(rowMenu).toHaveCount(0);
+    await expect(row).toBeFocused();
+  } finally {
+    await exit.release();
+  }
+});