From 59370d5dedcf48b2077c7c2a5146d9b92b30769b Mon Sep 17 00:00:00 2001
From: Sawyer Hood <sawyerjhood@gmail.com>
Date: Mon, 7 Sep 2026 07:51:30 -0700
Subject: [PATCH] Restore focus after controlled rename dialogs close

---
 .../components/dialogs/RenameDialog.test.tsx  | 43 +++++++++++++++++++
 .../src/components/dialogs/RenameDialog.tsx   | 34 ++++++++++++++-
 2 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 apps/app/src/components/dialogs/RenameDialog.test.tsx

diff --git a/apps/app/src/components/dialogs/RenameDialog.test.tsx b/apps/app/src/components/dialogs/RenameDialog.test.tsx
new file mode 100644
index 0000000000..89926651d5
--- /dev/null
+++ b/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));
+});
diff --git a/apps/app/src/components/dialogs/RenameDialog.tsx b/apps/app/src/components/dialogs/RenameDialog.tsx
index 31dff6372a..968e60a0d0 100644
--- a/apps/app/src/components/dialogs/RenameDialog.tsx
+++ b/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>