From 0941ea06f923d60d3e0d2364a72c49c3afbc8552 Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@umami.is>
Date: Thu, 30 Jul 2026 22:57:49 -0700
Subject: [PATCH] Wrap dialogs with ControlledDialog for consistent modal state
 management

---
 .../(main)/admin/teams/AdminTeamsTable.tsx    |  13 +-
 src/app/(main)/admin/users/UsersTable.tsx     |  21 +-
 .../users/[userId]/UserTwoFactorSettings.tsx  |  29 +--
 .../admin/websites/AdminWebsitesTable.tsx     |  13 +-
 .../boards/[boardId]/BoardEditColumn.tsx      |  49 ++---
 .../[websiteId]/ExpandedViewModal.tsx         |  43 +++--
 .../[websiteId]/replays/ReplayModal.tsx       |  43 +++--
 .../[websiteId]/sessions/SessionModal.tsx     |  49 ++---
 .../sessions/SessionProfileModal.tsx          |  33 ++--
 src/components/common/ControlledDialog.tsx    |  12 ++
 src/components/input/DateFilter.tsx           |  27 +--
 src/components/input/DialogButton.tsx         |  17 +-
 src/components/input/ReportEditButton.tsx     |  23 ++-
 .../modals/TwoFactorDisableModal.tsx          |  89 +++++----
 src/components/modals/TwoFactorSetupModal.tsx | 181 +++++++++---------
 .../modals/TwoFactorSuccessModal.tsx          |  83 ++++----
 16 files changed, 396 insertions(+), 329 deletions(-)
 create mode 100644 src/components/common/ControlledDialog.tsx

diff --git a/src/app/(main)/admin/teams/AdminTeamsTable.tsx b/src/app/(main)/admin/teams/AdminTeamsTable.tsx
index 697b19d25..2cc2b883a 100644
--- a/src/app/(main)/admin/teams/AdminTeamsTable.tsx
+++ b/src/app/(main)/admin/teams/AdminTeamsTable.tsx
@@ -1,6 +1,7 @@
 import { DataColumn, DataTable, Dialog, Icon, MenuItem, Modal, Row, Text } from '@umami/react-zen';
 import { useRouter } from 'next/navigation';
 import { useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { DateDistance } from '@/components/common/DateDistance';
 import Link from '@/components/common/Link';
 import { SortableLabel } from '@/components/common/SortableLabel';
@@ -93,11 +94,13 @@ export function AdminTeamsTable({
           </DataColumn>
         )}
       </DataTable>
-      <Modal isOpen={!!deleteTeam}>
-        <Dialog style={{ width: 400 }}>
-          <TeamDeleteForm teamId={deleteTeam} onClose={() => setDeleteTeam(null)} />
-        </Dialog>
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={!!deleteTeam}>
+          <Dialog style={{ width: 400 }}>
+            <TeamDeleteForm teamId={deleteTeam} onClose={() => setDeleteTeam(null)} />
+          </Dialog>
+        </Modal>
+      </ControlledDialog>
     </>
   );
 }
diff --git a/src/app/(main)/admin/users/UsersTable.tsx b/src/app/(main)/admin/users/UsersTable.tsx
index 4ec3604c0..5c3fdad74 100644
--- a/src/app/(main)/admin/users/UsersTable.tsx
+++ b/src/app/(main)/admin/users/UsersTable.tsx
@@ -1,6 +1,7 @@
 import { DataColumn, DataTable, Icon, MenuItem, Modal, Row, Text } from '@umami/react-zen';
 import { useRouter } from 'next/navigation';
 import { useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { DateDistance } from '@/components/common/DateDistance';
 import Link from '@/components/common/Link';
 import { SortableLabel } from '@/components/common/SortableLabel';
@@ -84,15 +85,17 @@ export function UsersTable({
           </DataColumn>
         )}
       </DataTable>
-      <Modal isOpen={!!deleteUser}>
-        <UserDeleteForm
-          userId={deleteUser?.id}
-          username={deleteUser?.username}
-          onClose={() => {
-            setDeleteUser(null);
-          }}
-        />
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={!!deleteUser}>
+          <UserDeleteForm
+            userId={deleteUser?.id}
+            username={deleteUser?.username}
+            onClose={() => {
+              setDeleteUser(null);
+            }}
+          />
+        </Modal>
+      </ControlledDialog>
     </>
   );
 }
diff --git a/src/app/(main)/admin/users/[userId]/UserTwoFactorSettings.tsx b/src/app/(main)/admin/users/[userId]/UserTwoFactorSettings.tsx
index c5160e542..80832d18e 100644
--- a/src/app/(main)/admin/users/[userId]/UserTwoFactorSettings.tsx
+++ b/src/app/(main)/admin/users/[userId]/UserTwoFactorSettings.tsx
@@ -13,6 +13,7 @@ import {
 } from '@umami/react-zen';
 import { useState } from 'react';
 import { Badge } from '@/components/common/Badge';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { TypeConfirmationForm } from '@/components/common/TypeConfirmationForm';
 import {
   useDeleteQuery,
@@ -112,19 +113,21 @@ export function UserTwoFactorSettings({ userId }: { userId: string }) {
         )}
       </Column>
 
-      <Modal isOpen={showReset} onOpenChange={open => !open && setShowReset(false)}>
-        <Dialog title={t(labels.twoFactorReset)} style={{ width: 400 }}>
-          <TypeConfirmationForm
-            confirmationValue={CONFIRM_VALUE}
-            onConfirm={handleReset}
-            onClose={() => setShowReset(false)}
-            isLoading={isResetting}
-            error={resetError}
-            buttonLabel={t(labels.reset)}
-            buttonVariant="danger"
-          />
-        </Dialog>
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={showReset} onOpenChange={open => !open && setShowReset(false)}>
+          <Dialog title={t(labels.twoFactorReset)} style={{ width: 400 }}>
+            <TypeConfirmationForm
+              confirmationValue={CONFIRM_VALUE}
+              onConfirm={handleReset}
+              onClose={() => setShowReset(false)}
+              isLoading={isResetting}
+              error={resetError}
+              buttonLabel={t(labels.reset)}
+              buttonVariant="danger"
+            />
+          </Dialog>
+        </Modal>
+      </ControlledDialog>
     </>
   );
 }
diff --git a/src/app/(main)/admin/websites/AdminWebsitesTable.tsx b/src/app/(main)/admin/websites/AdminWebsitesTable.tsx
index 3ae07d06b..0b8d0c2a7 100644
--- a/src/app/(main)/admin/websites/AdminWebsitesTable.tsx
+++ b/src/app/(main)/admin/websites/AdminWebsitesTable.tsx
@@ -2,6 +2,7 @@ import { DataColumn, DataTable, Dialog, Icon, MenuItem, Modal, Row, Text } from
 import { useRouter } from 'next/navigation';
 import { useState } from 'react';
 import { WebsiteDeleteForm } from '@/app/(main)/websites/[websiteId]/settings/WebsiteDeleteForm';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { DateDistance } from '@/components/common/DateDistance';
 import Link from '@/components/common/Link';
 import { SortableLabel } from '@/components/common/SortableLabel';
@@ -92,11 +93,13 @@ export function AdminWebsitesTable({ data = [], ...props }: { data: any[] }) {
           }}
         </DataColumn>
       </DataTable>
-      <Modal isOpen={!!deleteWebsite}>
-        <Dialog style={{ width: 400 }}>
-          <WebsiteDeleteForm websiteId={deleteWebsite} onClose={() => setDeleteWebsite(null)} />
-        </Dialog>
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={!!deleteWebsite}>
+          <Dialog style={{ width: 400 }}>
+            <WebsiteDeleteForm websiteId={deleteWebsite} onClose={() => setDeleteWebsite(null)} />
+          </Dialog>
+        </Modal>
+      </ControlledDialog>
     </>
   );
 }
diff --git a/src/app/(main)/boards/[boardId]/BoardEditColumn.tsx b/src/app/(main)/boards/[boardId]/BoardEditColumn.tsx
index cdc053a43..852e1660b 100644
--- a/src/app/(main)/boards/[boardId]/BoardEditColumn.tsx
+++ b/src/app/(main)/boards/[boardId]/BoardEditColumn.tsx
@@ -10,6 +10,7 @@ import {
   TooltipTrigger,
 } from '@umami/react-zen';
 import { useMemo, useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { Panel } from '@/components/common/Panel';
 import { useBoard, useMessages, useNavigation } from '@/components/hooks';
 import { Pencil, Plus, X } from '@/components/icons';
@@ -121,29 +122,31 @@ export function BoardEditColumn({
           </Box>
         </Column>
       )}
-      <Modal isOpen={showSelect} onOpenChange={setShowSelect}>
-        <Dialog
-          title={t(labels.selectComponent)}
-          style={{
-            width: '1200px',
-            maxWidth: 'calc(100vw - 40px)',
-            maxHeight: 'calc(100dvh - 40px)',
-            padding: '32px',
-          }}
-        >
-          {() => (
-            <BoardComponentSelect
-              teamId={teamId}
-              boardType={boardType}
-              boardEntityType={boardEntityType}
-              boardEntityId={boardEntityId}
-              initialConfig={component}
-              onSelect={handleSelect}
-              onClose={() => setShowSelect(false)}
-            />
-          )}
-        </Dialog>
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={showSelect} onOpenChange={setShowSelect}>
+          <Dialog
+            title={t(labels.selectComponent)}
+            style={{
+              width: '1200px',
+              maxWidth: 'calc(100vw - 40px)',
+              maxHeight: 'calc(100dvh - 40px)',
+              padding: '32px',
+            }}
+          >
+            {() => (
+              <BoardComponentSelect
+                teamId={teamId}
+                boardType={boardType}
+                boardEntityType={boardEntityType}
+                boardEntityId={boardEntityId}
+                initialConfig={component}
+                onSelect={handleSelect}
+                onClose={() => setShowSelect(false)}
+              />
+            )}
+          </Dialog>
+        </Modal>
+      </ControlledDialog>
     </Panel>
   );
 }
diff --git a/src/app/(main)/websites/[websiteId]/ExpandedViewModal.tsx b/src/app/(main)/websites/[websiteId]/ExpandedViewModal.tsx
index 3336d2654..f892aee1d 100644
--- a/src/app/(main)/websites/[websiteId]/ExpandedViewModal.tsx
+++ b/src/app/(main)/websites/[websiteId]/ExpandedViewModal.tsx
@@ -1,5 +1,6 @@
 import { Dialog, Modal } from '@umami/react-zen';
 import { WebsiteExpandedView } from '@/app/(main)/websites/[websiteId]/WebsiteExpandedView';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { useMobile, useNavigation } from '@/components/hooks';
 
 export function ExpandedViewModal({
@@ -28,25 +29,27 @@ export function ExpandedViewModal({
   };
 
   return (
-    <Modal isOpen={!!view} onOpenChange={handleOpenChange}>
-      <Dialog
-        style={{
-          maxWidth: 1320,
-          width: '100vw',
-          height: isMobile ? '100dvh' : 'calc(100dvh - 40px)',
-          overflow: 'hidden',
-        }}
-      >
-        {({ close }) => {
-          return (
-            <WebsiteExpandedView
-              websiteId={websiteId}
-              excludedIds={excludedIds}
-              onClose={() => handleClose(close)}
-            />
-          );
-        }}
-      </Dialog>
-    </Modal>
+    <ControlledDialog>
+      <Modal isOpen={!!view} onOpenChange={handleOpenChange}>
+        <Dialog
+          style={{
+            maxWidth: 1320,
+            width: '100vw',
+            height: isMobile ? '100dvh' : 'calc(100dvh - 40px)',
+            overflow: 'hidden',
+          }}
+        >
+          {({ close }) => {
+            return (
+              <WebsiteExpandedView
+                websiteId={websiteId}
+                excludedIds={excludedIds}
+                onClose={() => handleClose(close)}
+              />
+            );
+          }}
+        </Dialog>
+      </Modal>
+    </ControlledDialog>
   );
 }
diff --git a/src/app/(main)/websites/[websiteId]/replays/ReplayModal.tsx b/src/app/(main)/websites/[websiteId]/replays/ReplayModal.tsx
index 0e04fa5ea..004df69df 100644
--- a/src/app/(main)/websites/[websiteId]/replays/ReplayModal.tsx
+++ b/src/app/(main)/websites/[websiteId]/replays/ReplayModal.tsx
@@ -1,6 +1,7 @@
 'use client';
 import { Column, Dialog, Sheet, type SheetProps } from '@umami/react-zen';
 import { ReplayPlayback } from '@/app/(main)/websites/[websiteId]/replays/[replayId]/ReplayPlayback';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { useNavigation } from '@/components/hooks';
 import { buildPath } from '@/lib/url';
 import styles from './ReplayModal.module.css';
@@ -34,25 +35,27 @@ export function ReplayModal({ websiteId, replayId, className, ...props }: Replay
   };
 
   return (
-    <Sheet
-      {...props}
-      side="bottom"
-      size="calc(100dvh - 80px)"
-      className={modalClassName}
-      isOpen={!!activeReplayId}
-      onOpenChange={handleOpenChange}
-    >
-      <Column height="100%">
-        <Dialog className={styles.sheet}>
-          {({ close }) => (
-            <Column padding="6">
-              {activeReplayId && (
-                <ReplayPlayback websiteId={websiteId} replayId={activeReplayId} onClose={close} />
-              )}
-            </Column>
-          )}
-        </Dialog>
-      </Column>
-    </Sheet>
+    <ControlledDialog>
+      <Sheet
+        {...props}
+        side="bottom"
+        size="calc(100dvh - 80px)"
+        className={modalClassName}
+        isOpen={!!activeReplayId}
+        onOpenChange={handleOpenChange}
+      >
+        <Column height="100%">
+          <Dialog className={styles.sheet}>
+            {({ close }) => (
+              <Column padding="6">
+                {activeReplayId && (
+                  <ReplayPlayback websiteId={websiteId} replayId={activeReplayId} onClose={close} />
+                )}
+              </Column>
+            )}
+          </Dialog>
+        </Column>
+      </Sheet>
+    </ControlledDialog>
   );
 }
diff --git a/src/app/(main)/websites/[websiteId]/sessions/SessionModal.tsx b/src/app/(main)/websites/[websiteId]/sessions/SessionModal.tsx
index d9c41bc8a..83d87b7c4 100644
--- a/src/app/(main)/websites/[websiteId]/sessions/SessionModal.tsx
+++ b/src/app/(main)/websites/[websiteId]/sessions/SessionModal.tsx
@@ -1,6 +1,7 @@
 'use client';
 import { Column, Dialog, Sheet, type SheetProps } from '@umami/react-zen';
 import { SessionProfile } from '@/app/(main)/websites/[websiteId]/sessions/SessionProfile';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { useNavigation } from '@/components/hooks';
 import styles from './SessionModal.module.css';
 
@@ -23,28 +24,30 @@ export function SessionModal({ websiteId, className, ...props }: SessionModalPro
   };
 
   return (
-    <Sheet
-      side="bottom"
-      size="calc(100dvh - 80px)"
-      className={[styles.modal, className].filter(Boolean).join(' ')}
-      isOpen={!!session}
-      onOpenChange={handleOpenChange}
-      {...props}
-    >
-      <Column height="100%">
-        <Dialog className="rounded-lg">
-          {({ close }) => (
-            <Column padding="10">
-              <SessionProfile
-                websiteId={websiteId}
-                sessionId={session}
-                showReplays={!isSharePage}
-                onClose={() => close()}
-              />
-            </Column>
-          )}
-        </Dialog>
-      </Column>
-    </Sheet>
+    <ControlledDialog>
+      <Sheet
+        side="bottom"
+        size="calc(100dvh - 80px)"
+        className={[styles.modal, className].filter(Boolean).join(' ')}
+        isOpen={!!session}
+        onOpenChange={handleOpenChange}
+        {...props}
+      >
+        <Column height="100%">
+          <Dialog className="rounded-lg">
+            {({ close }) => (
+              <Column padding="10">
+                <SessionProfile
+                  websiteId={websiteId}
+                  sessionId={session}
+                  showReplays={!isSharePage}
+                  onClose={() => close()}
+                />
+              </Column>
+            )}
+          </Dialog>
+        </Column>
+      </Sheet>
+    </ControlledDialog>
   );
 }
diff --git a/src/app/(main)/websites/[websiteId]/sessions/SessionProfileModal.tsx b/src/app/(main)/websites/[websiteId]/sessions/SessionProfileModal.tsx
index 4de485c38..b2f74db2d 100644
--- a/src/app/(main)/websites/[websiteId]/sessions/SessionProfileModal.tsx
+++ b/src/app/(main)/websites/[websiteId]/sessions/SessionProfileModal.tsx
@@ -1,6 +1,7 @@
 'use client';
 import { Column, Dialog, Sheet } from '@umami/react-zen';
 import { useRouter, useSearchParams } from 'next/navigation';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { buildPath } from '@/lib/url';
 import styles from './SessionModal.module.css';
 import { SessionProfile } from './SessionProfile';
@@ -29,20 +30,22 @@ export function SessionProfileModal({
   };
 
   return (
-    <Sheet
-      side="bottom"
-      size="calc(100dvh - 80px)"
-      className={styles.modal}
-      isOpen
-      onOpenChange={handleOpenChange}
-    >
-      <Column height="100%">
-        <Dialog className="rounded-lg">
-          <Column padding="10">
-            <SessionProfile websiteId={websiteId} sessionId={sessionId} onClose={closeModal} />
-          </Column>
-        </Dialog>
-      </Column>
-    </Sheet>
+    <ControlledDialog>
+      <Sheet
+        side="bottom"
+        size="calc(100dvh - 80px)"
+        className={styles.modal}
+        isOpen
+        onOpenChange={handleOpenChange}
+      >
+        <Column height="100%">
+          <Dialog className="rounded-lg">
+            <Column padding="10">
+              <SessionProfile websiteId={websiteId} sessionId={sessionId} onClose={closeModal} />
+            </Column>
+          </Dialog>
+        </Column>
+      </Sheet>
+    </ControlledDialog>
   );
 }
diff --git a/src/components/common/ControlledDialog.tsx b/src/components/common/ControlledDialog.tsx
new file mode 100644
index 000000000..2030bd273
--- /dev/null
+++ b/src/components/common/ControlledDialog.tsx
@@ -0,0 +1,12 @@
+'use client';
+import { DialogTrigger } from '@umami/react-zen';
+import type { ReactElement } from 'react';
+
+export function ControlledDialog({ children }: { children: ReactElement }) {
+  return (
+    <DialogTrigger>
+      <span hidden />
+      {children}
+    </DialogTrigger>
+  );
+}
diff --git a/src/components/input/DateFilter.tsx b/src/components/input/DateFilter.tsx
index 84aad53c2..825bdac02 100644
--- a/src/components/input/DateFilter.tsx
+++ b/src/components/input/DateFilter.tsx
@@ -1,5 +1,6 @@
 import { Dialog, ListItem, ListSeparator, Modal, Select, type SelectProps } from '@umami/react-zen';
 import { Fragment, type Key, useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { DateDisplay } from '@/components/common/DateDisplay';
 import { useMessages, useMobile } from '@/components/hooks';
 import { DatePickerForm } from '@/components/metrics/DatePickerForm';
@@ -128,18 +129,20 @@ export function DateFilter({
         })}
       </Select>
       {showPicker && (
-        <Modal isOpen={true}>
-          <Dialog>
-            <DatePickerForm
-              startDate={startDate}
-              endDate={endDate}
-              minDate={new Date(2000, 0, 1)}
-              maxDate={getMaxSelectableDate()}
-              onChange={handlePickerChange}
-              onClose={() => setShowPicker(false)}
-            />
-          </Dialog>
-        </Modal>
+        <ControlledDialog>
+          <Modal isOpen={true}>
+            <Dialog>
+              <DatePickerForm
+                startDate={startDate}
+                endDate={endDate}
+                minDate={new Date(2000, 0, 1)}
+                maxDate={getMaxSelectableDate()}
+                onChange={handlePickerChange}
+                onClose={() => setShowPicker(false)}
+              />
+            </Dialog>
+          </Modal>
+        </ControlledDialog>
       )}
     </>
   );
diff --git a/src/components/input/DialogButton.tsx b/src/components/input/DialogButton.tsx
index 8776be6f5..aa81489f4 100644
--- a/src/components/input/DialogButton.tsx
+++ b/src/components/input/DialogButton.tsx
@@ -7,6 +7,7 @@ import {
   Modal,
 } from '@umami/react-zen';
 import type { CSSProperties, ReactNode } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { IconLabel } from '@/components/common/IconLabel';
 import { useMobile } from '@/components/hooks';
 
@@ -64,13 +65,15 @@ export function DialogButton({
 
   if (isOpen !== undefined) {
     return (
-      <Modal
-        isOpen={isOpen}
-        onOpenChange={onOpenChange}
-        placement={isMobile ? 'fullscreen' : 'center'}
-      >
-        {dialog}
-      </Modal>
+      <ControlledDialog>
+        <Modal
+          isOpen={isOpen}
+          onOpenChange={onOpenChange}
+          placement={isMobile ? 'fullscreen' : 'center'}
+        >
+          {dialog}
+        </Modal>
+      </ControlledDialog>
     );
   }
 
diff --git a/src/components/input/ReportEditButton.tsx b/src/components/input/ReportEditButton.tsx
index c335650ef..3d245b151 100644
--- a/src/components/input/ReportEditButton.tsx
+++ b/src/components/input/ReportEditButton.tsx
@@ -11,6 +11,7 @@ import {
   Text,
 } from '@umami/react-zen';
 import { type ReactNode, useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { useMessages } from '@/components/hooks';
 import { useDeleteQuery } from '@/components/hooks/queries/useDeleteQuery';
 import { Edit, MoreHorizontal, Trash } from '@/components/icons';
@@ -103,16 +104,18 @@ export function ReportEditButton({
       >
         {children}
       </DialogButton>
-      <Modal isOpen={showDelete} onOpenChange={open => !open && handleClose()}>
-        <AlertDialog
-          title={t(labels.delete)}
-          onConfirm={handleDelete}
-          onCancel={handleClose}
-          isDanger
-        >
-          <Row gap="1">{t(messages.confirmDelete, { target: name })}</Row>
-        </AlertDialog>
-      </Modal>
+      <ControlledDialog>
+        <Modal isOpen={showDelete} onOpenChange={open => !open && handleClose()}>
+          <AlertDialog
+            title={t(labels.delete)}
+            onConfirm={handleDelete}
+            onCancel={handleClose}
+            isDanger
+          >
+            <Row gap="1">{t(messages.confirmDelete, { target: name })}</Row>
+          </AlertDialog>
+        </Modal>
+      </ControlledDialog>
     </>
   );
 }
diff --git a/src/components/modals/TwoFactorDisableModal.tsx b/src/components/modals/TwoFactorDisableModal.tsx
index 859d2df75..20c7368af 100644
--- a/src/components/modals/TwoFactorDisableModal.tsx
+++ b/src/components/modals/TwoFactorDisableModal.tsx
@@ -1,7 +1,18 @@
 'use client';
-import { Button, Column, Dialog, Modal, Row, Text } from '@umami/react-zen';
-import { Form, FormButtons, FormField, FormSubmitButton, PasswordField } from '@umami/react-zen';
+import {
+  Button,
+  Column,
+  Dialog,
+  Form,
+  FormButtons,
+  FormField,
+  FormSubmitButton,
+  Modal,
+  PasswordField,
+  Text,
+} from '@umami/react-zen';
 import { useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { OtpInput } from '@/components/common/OtpInput';
 import { useMessages, useUpdateQuery } from '@/components/hooks';
 
@@ -24,41 +35,43 @@ export function TwoFactorDisableModal({ onClose, onSuccess }: TwoFactorDisableMo
   };
 
   return (
-    <Modal isOpen={true} onOpenChange={onClose}>
-      <Dialog title={t(labels.twoFactorDisableTitle)} style={{ width: 420 }}>
-        {({ close }) => (
-          <Column gap="4">
-            <Text>{t(messages.twoFactorDisableDescription)}</Text>
-            <Form onSubmit={handleSubmit} error={getErrorMessage(error)}>
-              <FormField label={t(labels.password)} name="password">
-                <PasswordField autoComplete="current-password" onChange={resetError} />
-              </FormField>
-              <Column gap="1">
-                <Text size="sm">{t(labels.twoFactorEnterCode)}</Text>
-                <OtpInput
-                  value={otpValue}
-                  onChange={val => {
-                    setOtpValue(val);
-                    resetError();
-                  }}
-                  disabled={isPending}
-                />
-              </Column>
-              <FormButtons>
-                <Button variant="outline" onPress={close} isDisabled={isPending}>
-                  {t(labels.cancel)}
-                </Button>
-                <FormSubmitButton
-                  variant="danger"
-                  isDisabled={isPending || otpValue.length !== 6 || !!error}
-                >
-                  {t(labels.twoFactorDisable)}
-                </FormSubmitButton>
-              </FormButtons>
-            </Form>
-          </Column>
-        )}
-      </Dialog>
-    </Modal>
+    <ControlledDialog>
+      <Modal isOpen={true} onOpenChange={onClose}>
+        <Dialog title={t(labels.twoFactorDisableTitle)} style={{ width: 420 }}>
+          {({ close }) => (
+            <Column gap="4">
+              <Text>{t(messages.twoFactorDisableDescription)}</Text>
+              <Form onSubmit={handleSubmit} error={getErrorMessage(error)}>
+                <FormField label={t(labels.password)} name="password">
+                  <PasswordField autoComplete="current-password" onChange={resetError} />
+                </FormField>
+                <Column gap="1">
+                  <Text size="sm">{t(labels.twoFactorEnterCode)}</Text>
+                  <OtpInput
+                    value={otpValue}
+                    onChange={val => {
+                      setOtpValue(val);
+                      resetError();
+                    }}
+                    disabled={isPending}
+                  />
+                </Column>
+                <FormButtons>
+                  <Button variant="outline" onPress={close} isDisabled={isPending}>
+                    {t(labels.cancel)}
+                  </Button>
+                  <FormSubmitButton
+                    variant="danger"
+                    isDisabled={isPending || otpValue.length !== 6 || !!error}
+                  >
+                    {t(labels.twoFactorDisable)}
+                  </FormSubmitButton>
+                </FormButtons>
+              </Form>
+            </Column>
+          )}
+        </Dialog>
+      </Modal>
+    </ControlledDialog>
   );
 }
diff --git a/src/components/modals/TwoFactorSetupModal.tsx b/src/components/modals/TwoFactorSetupModal.tsx
index 56c6b6457..75bd01f19 100644
--- a/src/components/modals/TwoFactorSetupModal.tsx
+++ b/src/components/modals/TwoFactorSetupModal.tsx
@@ -19,6 +19,7 @@ import {
 } from '@umami/react-zen';
 import { LucideCopy } from 'lucide-react';
 import { type ReactNode, useEffect, useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { OtpInput } from '@/components/common/OtpInput';
 import { useMessages, useUpdateQuery } from '@/components/hooks';
 import styles from './TwoFactorSetupModal.module.css';
@@ -123,96 +124,98 @@ export function TwoFactorSetupModal({ required, onClose }: TwoFactorSetupModalPr
   const preventDismiss = required ? () => {} : undefined;
 
   return (
-    <Modal isOpen={true} onOpenChange={preventDismiss}>
-      <Dialog title={t(labels.twoFactorSetupTitle)} className={styles.twoFactorSetupModal}>
-        {({ close }) => (
-          <Column gap="9">
-            <Text>
-              {required && t(messages.twoFactorSetupRequiredDescription) + ' '}
-              {t(messages.twoFactorSetupDescription)}
-            </Text>
+    <ControlledDialog>
+      <Modal isOpen={true} onOpenChange={preventDismiss}>
+        <Dialog title={t(labels.twoFactorSetupTitle)} className={styles.twoFactorSetupModal}>
+          {() => (
+            <Column gap="9">
+              <Text>
+                {required && `${t(messages.twoFactorSetupRequiredDescription)} `}
+                {t(messages.twoFactorSetupDescription)}
+              </Text>
 
-            {/* Step 1 - Scan QR Code */}
-            <Step
-              tag={t(labels.twoFactorStep1)}
-              title={t(labels.twoFactorScanQr)}
-              details={t(messages.twoFactorStep1Description)}
-            >
-              <Box padding="2" shadow="lg" borderRadius="lg" border width="fit">
-                <Row alignItems="center" gap="2">
-                  {qrCodeDataUrl && (
-                    <Image
-                      src={qrCodeDataUrl}
-                      alt="QR code"
-                      className={styles.qrCodeImage}
-                      borderRadius="lg"
-                    />
-                  )}
-                  <Column>
-                    <Text weight="bold">{t(labels.twoFactorCantScan)}</Text>
-                    <Column gap="2">
-                      <Text size="sm">{t(labels.twoFactorManualEntry)}</Text>
-                      <Code>{manualKey}</Code>
-                      <div>
-                        <Button variant="outline" onPress={handleCopy}>
-                          <Icon size="sm">
-                            <LucideCopy />
-                          </Icon>
-                          {copied ? t(labels.twoFactorCodeCopied) : t(labels.twoFactorCopyCode)}
-                        </Button>
-                      </div>
-                    </Column>
-                  </Column>
-                </Row>
-              </Box>
-            </Step>
-
-            {/* Step 2 - Enter Verification Code */}
-            <Step
-              tag={t(labels.twoFactorStep2)}
-              title={t(labels.twoFactorGetCode)}
-              details={t(messages.twoFactorStep2Description)}
-            >
-              <Column gap="3.5">
-                <Text size="sm" weight="bold">
-                  {t(labels.twoFactorEnterCode)}
-                </Text>
-                <OtpInput
-                  value={otpValue}
-                  onChange={val => {
-                    setOtpValue(val);
-                    if (error) setError(null);
-                  }}
-                  onComplete={handleConfirm}
-                  disabled={isConfirming}
-                />
-              </Column>
-            </Step>
-
-            {error && (
-              <Alert variant="danger">
-                <AlertTitle>{t(messages.error)}</AlertTitle>
-                <AlertDescription>{error}</AlertDescription>
-              </Alert>
-            )}
-
-            <Row gap="2" justifyContent="flex-end">
-              {!required && (
-                <Button variant="outline" onPress={handleCancel} isDisabled={isConfirming}>
-                  {t(labels.cancel)}
-                </Button>
-              )}
-              <Button
-                variant="primary"
-                onPress={() => handleConfirm()}
-                isDisabled={otpValue.length !== 6 || isConfirming || !!error}
+              {/* Step 1 - Scan QR Code */}
+              <Step
+                tag={t(labels.twoFactorStep1)}
+                title={t(labels.twoFactorScanQr)}
+                details={t(messages.twoFactorStep1Description)}
               >
-                {t(labels.confirm)}
-              </Button>
-            </Row>
-          </Column>
-        )}
-      </Dialog>
-    </Modal>
+                <Box padding="2" shadow="lg" borderRadius="lg" border width="fit">
+                  <Row alignItems="center" gap="2">
+                    {qrCodeDataUrl && (
+                      <Image
+                        src={qrCodeDataUrl}
+                        alt="QR code"
+                        className={styles.qrCodeImage}
+                        borderRadius="lg"
+                      />
+                    )}
+                    <Column>
+                      <Text weight="bold">{t(labels.twoFactorCantScan)}</Text>
+                      <Column gap="2">
+                        <Text size="sm">{t(labels.twoFactorManualEntry)}</Text>
+                        <Code>{manualKey}</Code>
+                        <div>
+                          <Button variant="outline" onPress={handleCopy}>
+                            <Icon size="sm">
+                              <LucideCopy />
+                            </Icon>
+                            {copied ? t(labels.twoFactorCodeCopied) : t(labels.twoFactorCopyCode)}
+                          </Button>
+                        </div>
+                      </Column>
+                    </Column>
+                  </Row>
+                </Box>
+              </Step>
+
+              {/* Step 2 - Enter Verification Code */}
+              <Step
+                tag={t(labels.twoFactorStep2)}
+                title={t(labels.twoFactorGetCode)}
+                details={t(messages.twoFactorStep2Description)}
+              >
+                <Column gap="3.5">
+                  <Text size="sm" weight="bold">
+                    {t(labels.twoFactorEnterCode)}
+                  </Text>
+                  <OtpInput
+                    value={otpValue}
+                    onChange={val => {
+                      setOtpValue(val);
+                      if (error) setError(null);
+                    }}
+                    onComplete={handleConfirm}
+                    disabled={isConfirming}
+                  />
+                </Column>
+              </Step>
+
+              {error && (
+                <Alert variant="danger">
+                  <AlertTitle>{t(messages.error)}</AlertTitle>
+                  <AlertDescription>{error}</AlertDescription>
+                </Alert>
+              )}
+
+              <Row gap="2" justifyContent="flex-end">
+                {!required && (
+                  <Button variant="outline" onPress={handleCancel} isDisabled={isConfirming}>
+                    {t(labels.cancel)}
+                  </Button>
+                )}
+                <Button
+                  variant="primary"
+                  onPress={() => handleConfirm()}
+                  isDisabled={otpValue.length !== 6 || isConfirming || !!error}
+                >
+                  {t(labels.confirm)}
+                </Button>
+              </Row>
+            </Column>
+          )}
+        </Dialog>
+      </Modal>
+    </ControlledDialog>
   );
 }
diff --git a/src/components/modals/TwoFactorSuccessModal.tsx b/src/components/modals/TwoFactorSuccessModal.tsx
index 3818087d7..2f7d775f0 100644
--- a/src/components/modals/TwoFactorSuccessModal.tsx
+++ b/src/components/modals/TwoFactorSuccessModal.tsx
@@ -1,6 +1,7 @@
 'use client';
 import { Button, Column, Dialog, Modal, Row, Text } from '@umami/react-zen';
 import { useState } from 'react';
+import { ControlledDialog } from '@/components/common/ControlledDialog';
 import { useMessages } from '@/components/hooks';
 import styles from './TwoFactorSuccessModal.module.css';
 
@@ -30,51 +31,53 @@ export function TwoFactorSuccessModal({ backupCodes, onClose }: TwoFactorSuccess
   };
 
   return (
-    <Modal isOpen={true}>
-      <Dialog title={t(labels.twoFactorSuccessTitle)} className={styles.dialog}>
-        {() => (
-          <Column gap="5">
-            <Text>{t(messages.twoFactorEnabledDescription)}</Text>
+    <ControlledDialog>
+      <Modal isOpen={true}>
+        <Dialog title={t(labels.twoFactorSuccessTitle)} className={styles.dialog}>
+          {() => (
+            <Column gap="5">
+              <Text>{t(messages.twoFactorEnabledDescription)}</Text>
 
-            <Column gap="2">
-              <Text weight="bold">{t(labels.twoFactorSaveBackupCodes)}</Text>
-              <Text>{t(messages.twoFactorBackupDescription)}</Text>
+              <Column gap="2">
+                <Text weight="bold">{t(labels.twoFactorSaveBackupCodes)}</Text>
+                <Text>{t(messages.twoFactorBackupDescription)}</Text>
 
-              <div className={styles.codesGrid}>
-                {backupCodes.map((code, i) => (
-                  <code key={i} className={styles.code}>
-                    {code}
-                  </code>
-                ))}
-              </div>
+                <div className={styles.codesGrid}>
+                  {backupCodes.map((code, i) => (
+                    <code key={i} className={styles.code}>
+                      {code}
+                    </code>
+                  ))}
+                </div>
 
-              <Row>
-                <Button variant="outline" onPress={handleDownload}>
-                  {t(labels.twoFactorDownloadCodes)}
+                <Row>
+                  <Button variant="outline" onPress={handleDownload}>
+                    {t(labels.twoFactorDownloadCodes)}
+                  </Button>
+                </Row>
+              </Column>
+
+              <Row gap="2" alignItems="center">
+                <input
+                  type="checkbox"
+                  id="saved-codes"
+                  checked={saved}
+                  onChange={e => setSaved(e.target.checked)}
+                />
+                <label htmlFor="saved-codes">
+                  <Text>{t(labels.twoFactorSavedCodes)}</Text>
+                </label>
+              </Row>
+
+              <Row justifyContent="flex-end">
+                <Button variant="primary" onPress={onClose} isDisabled={!saved}>
+                  {t(labels.close)}
                 </Button>
               </Row>
             </Column>
-
-            <Row gap="2" alignItems="center">
-              <input
-                type="checkbox"
-                id="saved-codes"
-                checked={saved}
-                onChange={e => setSaved(e.target.checked)}
-              />
-              <label htmlFor="saved-codes">
-                <Text>{t(labels.twoFactorSavedCodes)}</Text>
-              </label>
-            </Row>
-
-            <Row justifyContent="flex-end">
-              <Button variant="primary" onPress={onClose} isDisabled={!saved}>
-                {t(labels.close)}
-              </Button>
-            </Row>
-          </Column>
-        )}
-      </Dialog>
-    </Modal>
+          )}
+        </Dialog>
+      </Modal>
+    </ControlledDialog>
   );
 }