Do better

src/components/dialogs/EmailDialog/data/useAccountEmailState.ts

@@ -3,19 +3,40 @@ import {useQuery, useQueryClient} from '@tanstack/react-query'
 
 import {useAgent} from '#/state/session'
 
-export const isEmailVerifiedQueryKey = ['isEmailVerified'] as const
+export type AccountEmailState = {
+  isEmailVerified: boolean
+  email2FAEnabled: boolean
+}
 
-export function useInvalidateIsEmailVerified() {
+export const accountEmailStateQueryKey = ['isEmailVerified'] as const
+
+export function useInvalidateAccountEmailState() {
   const qc = useQueryClient()
 
   return useCallback(() => {
     return qc.invalidateQueries({
-      queryKey: isEmailVerifiedQueryKey,
+      queryKey: accountEmailStateQueryKey,
     })
   }, [qc])
 }
 
-export function useIsEmailVerified({
+export function useUpdateAccountEmailStateQueryCache() {
+  const qc = useQueryClient()
+
+  return useCallback(
+    (data: AccountEmailState) => {
+      return qc.setQueriesData(
+        {
+          queryKey: accountEmailStateQueryKey,
+        },
+        data,
+      )
+    },
+    [qc],
+  )
+}
+
+export function useAccountEmailState({
   onVerify,
 }: {
   onVerify?: () => void
@@ -24,16 +45,20 @@ export function useIsEmailVerified({
   const [prevIsEmailVerified, setPrevEmailIsVerified] = useState(
     !!agent.session?.emailConfirmed,
   )
-  const query = useQuery({
+  const query = useQuery<AccountEmailState>({
     enabled: !!agent.session,
-    initialData: {isEmailVerified: !!agent.session?.emailConfirmed},
+    initialData: {
+      isEmailVerified: !!agent.session?.emailConfirmed,
+      email2FAEnabled: !!agent.session?.emailAuthFactor,
+    },
     refetchOnWindowFocus: true,
-    queryKey: isEmailVerifiedQueryKey,
+    queryKey: accountEmailStateQueryKey,
     queryFn: async () => {
       // will also trigger updates to `#/state/session` data
       const {data} = await agent.resumeSession(agent.session!)
       return {
         isEmailVerified: !!data.emailConfirmed,
+        email2FAEnabled: !!data.emailAuthFactor,
       }
     },
   })

src/components/dialogs/EmailDialog/data/useConfirmEmail.ts

@@ -1,12 +1,13 @@
 import {useMutation} from '@tanstack/react-query'
 
 import {useAgent, useSession} from '#/state/session'
-import {useInvalidateIsEmailVerified} from '#/components/dialogs/EmailDialog/data/useIsEmailVerified'
+import {useUpdateAccountEmailStateQueryCache} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 
 export function useConfirmEmail() {
   const agent = useAgent()
   const {currentAccount} = useSession()
-  const invalidateIsEmailVerified = useInvalidateIsEmailVerified()
+  const updateAccountEmailStateQueryCache =
+    useUpdateAccountEmailStateQueryCache()
 
   return useMutation({
     mutationFn: async ({token}: {token: string}) => {
@@ -18,8 +19,11 @@ export function useConfirmEmail() {
         email: currentAccount.email,
         token: token.trim(),
       })
-      await agent.resumeSession(agent.session!)
-      await invalidateIsEmailVerified()
+      const {data} = await agent.resumeSession(agent.session!)
+      updateAccountEmailStateQueryCache({
+        isEmailVerified: !!data.emailConfirmed,
+        email2FAEnabled: !!data.emailAuthFactor,
+      })
     },
   })
 }

src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts

@@ -1,10 +1,13 @@
 import {useMutation} from '@tanstack/react-query'
 
 import {useAgent, useSession} from '#/state/session'
+import {useUpdateAccountEmailStateQueryCache} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 
 export function useManageEmail2FA() {
   const agent = useAgent()
   const {currentAccount} = useSession()
+  const updateAccountEmailStateQueryCache =
+    useUpdateAccountEmailStateQueryCache()
 
   return useMutation({
     mutationFn: async ({
@@ -22,7 +25,11 @@ export function useManageEmail2FA() {
         emailAuthFactor: enabled,
         token,
       })
-      await agent.resumeSession(agent.session!)
+      const {data} = await agent.resumeSession(agent.session!)
+      updateAccountEmailStateQueryCache({
+        isEmailVerified: !!data.emailConfirmed,
+        email2FAEnabled: !!data.emailAuthFactor,
+      })
     },
   })
 }

src/components/dialogs/EmailDialog/index.tsx

@@ -7,7 +7,7 @@ import {
   type StatefulControl,
   useStatefulDialogControl,
 } from '#/components/dialogs/Context'
-import {useIsEmailVerified} from '#/components/dialogs/EmailDialog/data/useIsEmailVerified'
+import {useAccountEmailState} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 import {Manage2FA} from '#/components/dialogs/EmailDialog/screens/Manage2FA'
 import {Update} from '#/components/dialogs/EmailDialog/screens/Update'
 import {VerificationReminder} from '#/components/dialogs/EmailDialog/screens/VerificationReminder'
@@ -24,7 +24,7 @@ export function useEmailDialogControl() {
 
 export function EmailDialog({control}: {control: StatefulControl<Screen>}) {
   const {_} = useLingui()
-  const {isEmailVerified} = useIsEmailVerified()
+  const {isEmailVerified} = useAccountEmailState()
   const onClose = useCallback(() => {
     if (!isEmailVerified) {
       if (control.value?.id === ScreenID.Verify) {

src/components/dialogs/EmailDialog/screens/Manage2FA/Disable.tsx

@@ -228,7 +228,7 @@ export function Disable() {
             {state.tokenStatus === 'pending' ? (
               <ButtonIcon icon={Loader} />
             ) : state.tokenStatus === 'success' ? (
-              <ButtonIcon icon={Loader} />
+              <ButtonIcon icon={Check} />
             ) : null}
           </Button>
         </>

src/components/dialogs/EmailDialog/screens/Manage2FA/index.tsx

@@ -1,8 +1,7 @@
 import {useState} from 'react'
 import {Trans} from '@lingui/macro'
 
-import {useSession} from '#/state/session'
-import {useIsEmailVerified} from '#/components/dialogs/EmailDialog/data/useIsEmailVerified'
+import {useAccountEmailState} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 import {Disable} from '#/components/dialogs/EmailDialog/screens/Manage2FA/Disable'
 import {Enable} from '#/components/dialogs/EmailDialog/screens/Manage2FA/Enable'
 import {
@@ -11,12 +10,10 @@ import {
 } from '#/components/dialogs/EmailDialog/types'
 
 export function Manage2FA({showScreen}: ScreenProps<ScreenID.Manage2FA>) {
+  const {isEmailVerified, email2FAEnabled} = useAccountEmailState()
   const [requestedAction, setRequestedAction] = useState<
     'enable' | 'disable' | null
   >(null)
-  // TODO confirm this is in sync
-  const {currentAccount} = useSession()
-  const {isEmailVerified} = useIsEmailVerified()
 
   if (!isEmailVerified) {
     showScreen({
@@ -32,24 +29,41 @@ export function Manage2FA({showScreen}: ScreenProps<ScreenID.Manage2FA>) {
         })
       },
     })
+
     return null
   }
 
-  if (currentAccount?.emailAuthFactor && requestedAction !== 'enable') {
+  /*
+   * Wacky state handling so that once 2FA settings change, we don't show the
+   * wrong step of this form - esb
+   */
+
+  if (email2FAEnabled) {
     if (!requestedAction) {
       setRequestedAction('disable')
+      return <Disable />
     }
-    return <Disable />
-  } else if (
-    !currentAccount?.emailAuthFactor &&
-    requestedAction !== 'disable'
-  ) {
+
+    if (requestedAction === 'disable') {
+      return <Disable />
+    }
+    if (requestedAction === 'enable') {
+      return <Enable />
+    }
+  } else {
     if (!requestedAction) {
       setRequestedAction('enable')
+      return <Enable />
     }
 
-    return <Enable />
+    if (requestedAction === 'disable') {
+      return <Disable />
+    }
+    if (requestedAction === 'enable') {
+      return <Enable />
+    }
   }
 
+  // should never happen
   return null
 }

src/components/dialogs/EmailDialog/screens/Verify.tsx

@@ -11,8 +11,8 @@ import {Admonition} from '#/components/Admonition'
 import {Button, ButtonIcon, ButtonText} from '#/components/Button'
 import {ResendEmailText} from '#/components/dialogs/EmailDialog/components/ResendEmailText'
 import {TokenField} from '#/components/dialogs/EmailDialog/components/TokenField'
+import {useAccountEmailState} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 import {useConfirmEmail} from '#/components/dialogs/EmailDialog/data/useConfirmEmail'
-import {useIsEmailVerified} from '#/components/dialogs/EmailDialog/data/useIsEmailVerified'
 import {useRequestEmailVerification} from '#/components/dialogs/EmailDialog/data/useRequestEmailVerification'
 import {
   type ScreenID,
@@ -108,7 +108,7 @@ export function Verify({config}: ScreenProps<ScreenID.Verify>) {
       })
     }
   }, [config, dispatch])
-  useIsEmailVerified({onVerify})
+  useAccountEmailState({onVerify})
 
   const handleRequestEmailVerification = async () => {
     dispatch({