fix(cdk/dialog): Allow the dialog to recapture focus after active ele…

…ment 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

@@ -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();

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],