Restore prior focus after programmatic rename dialogs close

apps/app/src/components/dialogs/RenameDialog.test.tsx

@@ -0,0 +1,43 @@
+// @vitest-environment jsdom
+
+import { cleanup, render, screen, waitFor } from "@testing-library/react";
+import { afterEach, expect, it, vi } from "vitest";
+import { RenameDialog, RenameDialogContent } from "./RenameDialog";
+
+function RenameDialogHarness({ open }: { open: boolean }) {
+  return (
+    <>
+      <textarea aria-label="Composer" />
+      <RenameDialog open={open} onOpenChange={vi.fn()}>
+        {(inputRef) => (
+          <RenameDialogContent
+            entityLabel="thread"
+            initialName="Current title"
+            pending={false}
+            autoCapitalize="sentences"
+            onRename={vi.fn()}
+            inputRef={inputRef}
+          />
+        )}
+      </RenameDialog>
+    </>
+  );
+}
+
+afterEach(cleanup);
+
+it("restores focus after a programmatically opened rename dialog closes", async () => {
+  const view = render(<RenameDialogHarness open={false} />);
+  const composer = screen.getByRole("textbox", { name: "Composer" });
+  composer.focus();
+
+  view.rerender(<RenameDialogHarness open />);
+  await waitFor(() =>
+    expect(document.activeElement).toBe(
+      screen.getByRole("textbox", { name: "Thread name" }),
+    ),
+  );
+
+  view.rerender(<RenameDialogHarness open={false} />);
+  await waitFor(() => expect(document.activeElement).toBe(composer));
+});

apps/app/src/components/dialogs/RenameDialog.tsx

@@ -1,6 +1,8 @@
 import { capitalize } from "@bb/thread-view";
 import {
+  useCallback,
   useId,
+  useRef,
   useState,
   type FormEvent,
   type ReactNode,
@@ -32,12 +34,42 @@ export function RenameDialog({
   shellClassName,
   children,
 }: RenameDialogProps) {
-  const { inputRef, handleOpenAutoFocus } = useRenameDialogAutoFocus();
+  const { inputRef, handleOpenAutoFocus: focusInputOnOpen } =
+    useRenameDialogAutoFocus();
+  const returnFocusRef = useRef<HTMLElement | null>(null);
+  const capturedOpenFocusRef = useRef(false);
+  const handleOpenAutoFocus = useCallback(
+    (event: Event) => {
+      if (!capturedOpenFocusRef.current) {
+        const activeElement = document.activeElement;
+        returnFocusRef.current =
+          activeElement instanceof HTMLElement &&
+          activeElement !== document.body
+            ? activeElement
+            : null;
+        capturedOpenFocusRef.current = true;
+      }
+      focusInputOnOpen(event);
+    },
+    [focusInputOnOpen],
+  );
+  const handleAfterCloseAutoFocus = useCallback(() => {
+    const returnFocus = returnFocusRef.current;
+    returnFocusRef.current = null;
+    capturedOpenFocusRef.current = false;
+    if (
+      returnFocus?.isConnected &&
+      returnFocus.closest('[aria-hidden="true"], [inert]') === null
+    ) {
+      returnFocus.focus({ preventScroll: true });
+    }
+  }, []);
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
       <DialogContent
         className={shellClassName}
         onOpenAutoFocus={handleOpenAutoFocus}
+        onAfterCloseAutoFocus={handleAfterCloseAutoFocus}
       >
         {children(inputRef)}
       </DialogContent>