[PM-24178] Handle dialog focus when launched from a menu item (#18208)

libs/components/src/dialog/dialog.service.ts

@@ -10,7 +10,7 @@ import { ComponentPortal, Portal } from "@angular/cdk/portal";
 import { Injectable, Injector, TemplateRef, inject } from "@angular/core";
 import { takeUntilDestroyed } from "@angular/core/rxjs-interop";
 import { NavigationEnd, Router } from "@angular/router";
-import { filter, firstValueFrom, map, Observable, Subject, switchMap } from "rxjs";
+import { filter, firstValueFrom, map, Observable, Subject, switchMap, take } from "rxjs";
 
 import { AuthService } from "@bitwarden/common/auth/abstractions/auth.service";
 import { AuthenticationStatus } from "@bitwarden/common/auth/enums/authentication-status";
@@ -62,7 +62,7 @@ export abstract class DialogRef<R = unknown, C = unknown> implements Pick<
 
 export type DialogConfig<D = unknown, R = unknown> = Pick<
   CdkDialogConfig<D, R>,
-  "data" | "disableClose" | "ariaModal" | "positionStrategy" | "height" | "width"
+  "data" | "disableClose" | "ariaModal" | "positionStrategy" | "height" | "width" | "restoreFocus"
 >;
 
 /**
@@ -242,6 +242,11 @@ export class DialogService {
     };
 
     ref.cdkDialogRefBase = this.dialog.open<R, D, C>(componentOrTemplateRef, _config);
+
+    if (config?.restoreFocus === undefined) {
+      this.setRestoreFocusEl<R, C>(ref);
+    }
+
     return ref;
   }
 
@@ -305,6 +310,48 @@ export class DialogService {
     return this.activeDrawer?.close();
   }
 
+  /**
+   * Configure the dialog to return focus to the previous active element upon closing.
+   * @param ref CdkDialogRef
+   *
+   * The cdk dialog already has the optional directive `cdkTrapFocusAutoCapture` to capture the
+   * current active element and return focus to it upon close. However, it does not have a way to
+   * delay the capture of the element. We need this delay in some situations, where the active
+   * element may be changing as the dialog is opening, and we want to wait for that to settle.
+   *
+   * For example -- the menu component often contains menu items that open dialogs. When the dialog
+   * opens, the menu is closing and is setting focus back to the menu trigger since the menu item no
+   * longer exists. We want to capture the menu trigger as the active element, not the about-to-be-
+   * nonexistent menu item. If we wait a tick, we can let the menu finish that focus move.
+   */
+  private setRestoreFocusEl<R = unknown, C = unknown>(ref: CdkDialogRef<R, C>) {
+    /**
+     * First, capture the current active el with no delay so that we can support normal use cases
+     * where we are not doing manual focus management
+     */
+    const activeEl = document.activeElement;
+
+    const restoreFocusTimeout = setTimeout(() => {
+      let restoreFocusEl = activeEl;
+
+      /**
+       * If the original active element is no longer connected, it's because we purposely removed it
+       * from the DOM and have moved focus. Select the new active element instead.
+       */
+      if (!restoreFocusEl?.isConnected) {
+        restoreFocusEl = document.activeElement;
+      }
+
+      if (restoreFocusEl instanceof HTMLElement) {
+        ref.cdkDialogRefBase.config.restoreFocus = restoreFocusEl;
+      }
+    }, 0);
+
+    ref.closed.pipe(take(1)).subscribe(() => {
+      clearTimeout(restoreFocusTimeout);
+    });
+  }
+
   /** The injector that is passed to the opened dialog */
   private createInjector(opts: { data: unknown; dialogRef: DialogRef }): Injector {
     return Injector.create({

libs/components/src/dialog/dialog/dialog.component.html

@@ -6,7 +6,6 @@
     isDrawer ? 'tw-h-full tw-border-t-0' : 'tw-rounded-t-xl md:tw-rounded-xl tw-shadow-lg',
   ]"
   cdkTrapFocus
-  cdkTrapFocusAutoCapture
 >
   @let showHeaderBorder = bodyHasScrolledFrom().top;
   <header
@@ -23,8 +22,8 @@
       bitTypography="h3"
       noMargin
       class="tw-text-main tw-mb-0 tw-line-clamp-2 tw-text-ellipsis tw-break-words focus-visible:tw-outline-none"
-      cdkFocusInitial
       tabindex="-1"
+      #dialogHeader
     >
       {{ title() }}
       @if (subtitle(); as subtitleText) {

libs/components/src/dialog/dialog/dialog.component.ts

@@ -11,6 +11,7 @@ import {
   DestroyRef,
   computed,
   signal,
+  AfterViewInit,
 } from "@angular/core";
 import { toObservable } from "@angular/core/rxjs-interop";
 import { combineLatest, switchMap } from "rxjs";
@@ -62,8 +63,10 @@ const drawerSizeToWidth = {
     SpinnerComponent,
   ],
 })
-export class DialogComponent {
+export class DialogComponent implements AfterViewInit {
   private readonly destroyRef = inject(DestroyRef);
+  private readonly dialogHeader =
+    viewChild.required<ElementRef<HTMLHeadingElement>>("dialogHeader");
   private readonly scrollableBody = viewChild.required(CdkScrollable);
   private readonly scrollBottom = viewChild.required<ElementRef<HTMLDivElement>>("scrollBottom");
 
@@ -141,6 +144,22 @@ export class DialogComponent {
     return [...baseClasses, this.width(), ...sizeClasses, ...animationClasses];
   });
 
+  ngAfterViewInit() {
+    /**
+     * Wait a tick for any focus management to occur on the trigger element before moving focus to
+     * the dialog header. We choose the dialog header because it is always present, unlike possible
+     * interactive elements.
+     *
+     * We are doing this manually instead of using `cdkTrapFocusAutoCapture` and `cdkFocusInitial`
+     * because we need this delay behavior.
+     */
+    const headerFocusTimeout = setTimeout(() => {
+      this.dialogHeader().nativeElement.focus();
+    }, 0);
+
+    this.destroyRef.onDestroy(() => clearTimeout(headerFocusTimeout));
+  }
+
   handleEsc(event: Event) {
     if (!this.dialogRef?.disableClose) {
       this.dialogRef?.close();

libs/components/src/menu/menu-trigger-for.directive.ts

@@ -192,7 +192,7 @@ export class MenuTriggerForDirective implements OnDestroy {
       return;
     }
 
-    const escKey = this.overlayRef.keydownEvents().pipe(
+    const keyEvents = this.overlayRef.keydownEvents().pipe(
       filter((event: KeyboardEvent) => {
         const keys = this.menu().ariaRole() === "menu" ? ["Escape", "Tab"] : ["Escape"];
         return keys.includes(event.key);
@@ -202,8 +202,8 @@ export class MenuTriggerForDirective implements OnDestroy {
     const detachments = this.overlayRef.detachments();
 
     const closeEvents = isContextMenu
-      ? merge(detachments, escKey, menuClosed)
-      : merge(detachments, escKey, this.overlayRef.backdropClick(), menuClosed);
+      ? merge(detachments, keyEvents, menuClosed)
+      : merge(detachments, keyEvents, this.overlayRef.backdropClick(), menuClosed);
 
     this.closedEventsSub = closeEvents
       .pipe(takeUntil(this.overlayRef.detachments()))
@@ -215,9 +215,9 @@ export class MenuTriggerForDirective implements OnDestroy {
           event.preventDefault();
         }
 
-        if (event instanceof KeyboardEvent && (event.key === "Tab" || event.key === "Escape")) {
-          this.elementRef.nativeElement.focus();
-        }
+        // Move focus to the menu trigger, since any active menu items are about to be destroyed
+        this.elementRef.nativeElement.focus();
+
         this.destroyMenu();
       });
   }

libs/components/src/tooltip/tooltip.directive.ts

@@ -11,6 +11,7 @@ import {
   signal,
   model,
   computed,
+  OnDestroy,
 } from "@angular/core";
 
 import { TooltipPositionIdentifier, tooltipPositions } from "./tooltip-positions";
@@ -32,7 +33,7 @@ export const TOOLTIP_DELAY_MS = 800;
     "[attr.aria-describedby]": "resolvedDescribedByIds()",
   },
 })
-export class TooltipDirective implements OnInit {
+export class TooltipDirective implements OnInit, OnDestroy {
   private static nextId = 0;
   /**
    * The value of this input is forwarded to the tooltip.component to render
@@ -51,6 +52,7 @@ export class TooltipDirective implements OnInit {
 
   private readonly isVisible = signal(false);
   private overlayRef: OverlayRef | undefined;
+  private showTimeoutId: ReturnType<typeof setTimeout> | undefined;
   private elementRef = inject<ElementRef<HTMLElement>>(ElementRef);
   private overlay = inject(Overlay);
   private viewContainerRef = inject(ViewContainerRef);
@@ -81,13 +83,29 @@ export class TooltipDirective implements OnInit {
     }),
   );
 
+  /**
+   * Clear any pending show timeout
+   *
+   * Use cases: prevent tooltip from appearing after hide; clear existing timeout before showing a
+   * new tooltip
+   */
+  private clearTimeout() {
+    if (this.showTimeoutId !== undefined) {
+      clearTimeout(this.showTimeoutId);
+      this.showTimeoutId = undefined;
+    }
+  }
+
   private destroyTooltip = () => {
+    this.clearTimeout();
     this.overlayRef?.dispose();
     this.overlayRef = undefined;
     this.isVisible.set(false);
   };
 
   protected showTooltip = () => {
+    this.clearTimeout();
+
     if (!this.overlayRef) {
       this.overlayRef = this.overlay.create({
         ...this.defaultPopoverConfig,
@@ -97,8 +115,9 @@ export class TooltipDirective implements OnInit {
       this.overlayRef.attach(this.tooltipPortal);
     }
 
-    setTimeout(() => {
+    this.showTimeoutId = setTimeout(() => {
       this.isVisible.set(true);
+      this.showTimeoutId = undefined;
     }, TOOLTIP_DELAY_MS);
   };
 
@@ -134,4 +153,8 @@ export class TooltipDirective implements OnInit {
   ngOnInit() {
     this.positionStrategy.withPositions(this.computePositions(this.tooltipPosition()));
   }
+
+  ngOnDestroy(): void {
+    this.destroyTooltip();
+  }
 }

libs/components/src/tooltip/tooltip.spec.ts

@@ -41,6 +41,7 @@ interface OverlayLike {
 interface OverlayRefStub {
   attach: (portal: ComponentPortal<unknown>) => unknown;
   updatePosition: () => void;
+  dispose: () => void;
 }
 
 describe("TooltipDirective (visibility only)", () => {
@@ -68,6 +69,7 @@ describe("TooltipDirective (visibility only)", () => {
         },
       })),
       updatePosition: jest.fn(),
+      dispose: jest.fn(),
     };
 
     const overlayMock: OverlayLike = {