From 7c1b1466406d3d5bda66328379d8728ba3aaaa39 Mon Sep 17 00:00:00 2001
From: Miles Malerba <mmalerba@google.com>
Date: Fri, 10 May 2024 15:48:12 -0700
Subject: [PATCH] fix(cdk/dialog): Allow the dialog to recapture focus after
 active element removal

I initially implemented this functionality in the focus-trap, though I'm
not sure we necissarily want it in all focus trapping scenarios. We may
want to add a similar configurable `autofocus` capability that the
dialog has to focus-trap. For now I'm applying the fix to dialog and
we'll consider adding it to the focus-trap more generally as a followup
FR at some point in the future.
---
 src/cdk/dialog/dialog-container.ts | 15 ++++++++++++
 src/material/dialog/dialog.spec.ts | 38 ++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/src/cdk/dialog/dialog-container.ts b/src/cdk/dialog/dialog-container.ts
index 7e135978c3f5..95fb35519fb9 100644
--- a/src/cdk/dialog/dialog-container.ts
+++ b/src/cdk/dialog/dialog-container.ts
@@ -107,6 +107,16 @@ export class CdkDialogContainer<C extends DialogConfig = DialogConfig>
 
   private _injector = inject(Injector);
 
+  private _recaptureOnFocusout = () => {
+    // One reason focus could leave the dialog is if the currently focused element was removed from
+    // the DOM. In that case we will get a focusout event immediately prior to the element's
+    // removal. We delay the recapture by a microtask so that we don't try to move focus back into
+    // the dialog by focusing the same element whose removal triggered the event.
+    queueMicrotask(() => {
+      this._recaptureFocus();
+    });
+  };
+
   constructor(
     protected _elementRef: ElementRef,
     protected _focusTrapFactory: FocusTrapFactory,
@@ -141,6 +151,9 @@ export class CdkDialogContainer<C extends DialogConfig = DialogConfig>
   }
 
   protected _contentAttached() {
+    this._ngZone.runOutsideAngular(() => {
+      this._elementRef.nativeElement.addEventListener('focusout', this._recaptureOnFocusout);
+    });
     this._initializeFocusTrap();
     this._handleBackdropClicks();
     this._captureInitialFocus();
@@ -155,6 +168,7 @@ export class CdkDialogContainer<C extends DialogConfig = DialogConfig>
   }
 
   ngOnDestroy() {
+    this._elementRef.nativeElement.removeEventListener('focusout', this._recaptureOnFocusout);
     this._restoreFocus();
   }
 
@@ -386,6 +400,7 @@ export class CdkDialogContainer<C extends DialogConfig = DialogConfig>
   private _handleBackdropClicks() {
     // Clicking on the backdrop will move focus out of dialog.
     // Recapture it if closing via the backdrop is disabled.
+    // TODO(mmalerba): This may not be necessary now that we recapture on focusout.
     this._overlayRef.backdropClick().subscribe(() => {
       if (this._config.disableClose) {
         this._recaptureFocus();
diff --git a/src/material/dialog/dialog.spec.ts b/src/material/dialog/dialog.spec.ts
index efa1d782d0fa..342029133721 100644
--- a/src/material/dialog/dialog.spec.ts
+++ b/src/material/dialog/dialog.spec.ts
@@ -1129,6 +1129,32 @@ describe('MDC-based MatDialog', () => {
     }),
   );
 
+  fit('should recapture focus when the focused element is removed from the DOM', fakeAsync(async () => {
+    TestBed.inject(NgZone).run(() =>
+      dialog.open(ContentWithConditionalButton, {
+        disableClose: true,
+        viewContainerRef: testViewContainerRef,
+        autoFocus: 'button',
+      }),
+    );
+
+    viewContainerFixture.detectChanges();
+    flush();
+
+    const cancelButton = document.querySelector<HTMLElement>('.cancel')!;
+    const okButton = document.querySelector<HTMLElement>('.ok');
+
+    expect(document.activeElement).toEqual(cancelButton);
+
+    TestBed.inject(NgZone).run(() => {
+      cancelButton.click();
+    });
+    viewContainerFixture.detectChanges();
+    flush();
+
+    expect(document.activeElement).toEqual(okButton);
+  }));
+
   describe('hasBackdrop option', () => {
     it('should have a backdrop', () => {
       dialog.open(PizzaMsg, {hasBackdrop: true, viewContainerRef: testViewContainerRef});
@@ -2394,6 +2420,18 @@ class ModuleBoundDialogChildComponent {
   constructor(public service: ModuleBoundDialogService) {}
 }
 
+@Component({
+  template: `
+    @if (!disableCancel) {
+      <button class="cancel" (click)="disableCancel = true">Cancel</button>
+    }
+    <button class="ok">Ok</button>
+  `,
+})
+class ContentWithConditionalButton {
+  disableCancel = false;
+}
+
 @NgModule({
   imports: [ModuleBoundDialogComponent, ModuleBoundDialogChildComponent],
   providers: [ModuleBoundDialogService],