Return focus after dialogs close (#11999)

src/dialogs/make-dialog-manager.ts

@@ -1,6 +1,9 @@
 import { HASSDomEvent, ValidHassDomEvent } from "../common/dom/fire_event";
 import { mainWindow } from "../common/dom/get_main_window";
 import { ProvideHassElement } from "../mixins/provide-hass-lit-mixin";
+import { ancestorsWithProperty } from "../common/dom/ancestors-with-property";
+import { deepActiveElement } from "../common/dom/deep-active-element";
+import { nextRender } from "../common/util/render-status";
 
 declare global {
   // for fire event
@@ -40,7 +43,17 @@ export interface DialogState {
   dialogParams?: unknown;
 }
 
-const LOADED = {};
+interface LoadedDialogInfo {
+  element: Promise<HassDialog>;
+  closedFocusTargets?: Set<Element>;
+}
+
+interface LoadedDialogsDict {
+  [tag: string]: LoadedDialogInfo;
+}
+
+const LOADED: LoadedDialogsDict = {};
+export const FOCUS_TARGET = Symbol.for("HA focus target");
 
 export const showDialog = async (
   element: HTMLElement & ProvideHassElement,
@@ -60,11 +73,24 @@ export const showDialog = async (
       }
       return;
     }
-    LOADED[dialogTag] = dialogImport().then(() => {
-      const dialogEl = document.createElement(dialogTag) as HassDialog;
-      element.provideHass(dialogEl);
-      return dialogEl;
-    });
+    LOADED[dialogTag] = {
+      element: dialogImport().then(() => {
+        const dialogEl = document.createElement(dialogTag) as HassDialog;
+        element.provideHass(dialogEl);
+        return dialogEl;
+      }),
+    };
+  }
+
+  // Get the focus targets after the dialog closes, but keep the original if dialog is being replaced
+  if (mainWindow.history.state?.replaced) {
+    LOADED[dialogTag].closedFocusTargets =
+      LOADED[mainWindow.history.state.dialog].closedFocusTargets;
+  } else {
+    LOADED[dialogTag].closedFocusTargets = ancestorsWithProperty(
+      deepActiveElement(),
+      FOCUS_TARGET
+    );
   }
 
   if (addHistory) {
@@ -93,25 +119,29 @@ export const showDialog = async (
       );
     }
   }
-  const dialogElement = await LOADED[dialogTag];
+
+  const dialogElement = await LOADED[dialogTag].element;
+  dialogElement.addEventListener("dialog-closed", _handleClosedFocus);
+
   // Append it again so it's the last element in the root,
   // so it's guaranteed to be on top of the other elements
   root.appendChild(dialogElement);
   dialogElement.showDialog(dialogParams);
 };
 
-export const replaceDialog = () => {
+export const replaceDialog = (dialogElement: HassDialog) => {
   mainWindow.history.replaceState(
     { ...mainWindow.history.state, replaced: true },
     ""
   );
+  dialogElement.removeEventListener("dialog-closed", _handleClosedFocus);
 };
 
 export const closeDialog = async (dialogTag: string): Promise<boolean> => {
   if (!(dialogTag in LOADED)) {
     return true;
   }
-  const dialogElement: HassDialog = await LOADED[dialogTag];
+  const dialogElement = await LOADED[dialogTag].element;
   if (dialogElement.closeDialog) {
     return dialogElement.closeDialog() !== false;
   }
@@ -137,3 +167,33 @@ export const makeDialogManager = (
     }
   );
 };
+
+const _handleClosedFocus = async (ev: HASSDomEvent<DialogClosedParams>) => {
+  const closedFocusTargets = LOADED[ev.detail.dialog].closedFocusTargets;
+  delete LOADED[ev.detail.dialog].closedFocusTargets;
+  if (!closedFocusTargets) return;
+
+  // Undo whatever the browser focused to provide easy checking
+  let focusedElement = deepActiveElement();
+  if (focusedElement instanceof HTMLElement) focusedElement.blur();
+
+  // Make sure backdrop is fully updated before trying (especially needed for underlay dialogs)
+  await nextRender();
+
+  // Try all targets in order and stop when one works
+  for (const focusTarget of closedFocusTargets) {
+    if (focusTarget instanceof HTMLElement) {
+      focusTarget.focus();
+      focusedElement = deepActiveElement();
+      if (focusedElement && focusedElement !== document.body) return;
+    }
+  }
+
+  if (__DEV__) {
+    // eslint-disable-next-line no-console
+    console.warn(
+      "Failed to focus any targets after closing dialog: %o",
+      closedFocusTargets
+    );
+  }
+};