fix email not updating in session state after email change

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

@@ -1,6 +1,6 @@
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent, useSession} from '#/state/session'
+import {useAgent, useSession, useSessionApi} from '#/state/session'
 
 export function useConfirmEmail({
   onSuccess,
@@ -8,6 +8,7 @@ export function useConfirmEmail({
 }: {onSuccess?: () => void; onError?: () => void} = {}) {
   const agent = useAgent()
   const {currentAccount} = useSession()
+  const {partialRefreshSession} = useSessionApi()
 
   return useMutation({
     mutationFn: async ({token}: {token: string}) => {
@@ -19,8 +20,7 @@ export function useConfirmEmail({
         email: currentAccount.email.trim(),
         token: token.trim(),
       })
-      // will update session state at root of app
-      await agent.resumeSession(agent.session!)
+      await partialRefreshSession()
     },
     onSuccess,
     onError,

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

@@ -1,10 +1,11 @@
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent, useSession} from '#/state/session'
+import {useAgent, useSession, useSessionApi} from '#/state/session'
 
 export function useManageEmail2FA() {
   const agent = useAgent()
   const {currentAccount} = useSession()
+  const {partialRefreshSession} = useSessionApi()
 
   return useMutation({
     mutationFn: async ({
@@ -22,8 +23,7 @@ export function useManageEmail2FA() {
         emailAuthFactor: enabled,
         token,
       })
-      // will update session state at root of app
-      await agent.resumeSession(agent.session!)
+      await partialRefreshSession()
     },
   })
 }

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

@@ -1,19 +1,21 @@
 import {useMutation} from '@tanstack/react-query'
 
-import {useAgent} from '#/state/session'
+import {useAgent, useSessionApi} from '#/state/session'
 import {useRequestEmailUpdate} from '#/components/dialogs/EmailDialog/data/useRequestEmailUpdate'
 
 async function updateEmailAndRefreshSession(
   agent: ReturnType<typeof useAgent>,
+  partialRefreshSession: () => Promise<void>,
   email: string,
   token?: string,
 ) {
   await agent.com.atproto.server.updateEmail({email: email.trim(), token})
-  await agent.resumeSession(agent.session!)
+  await partialRefreshSession()
 }
 
 export function useUpdateEmail() {
   const agent = useAgent()
+  const {partialRefreshSession} = useSessionApi()
   const {mutateAsync: requestEmailUpdate} = useRequestEmailUpdate()
 
   return useMutation<
@@ -23,7 +25,12 @@ export function useUpdateEmail() {
   >({
     mutationFn: async ({email, token}: {email: string; token?: string}) => {
       if (token) {
-        await updateEmailAndRefreshSession(agent, email, token)
+        await updateEmailAndRefreshSession(
+          agent,
+          partialRefreshSession,
+          email,
+          token,
+        )
         return {
           status: 'success',
         }
@@ -34,7 +41,12 @@ export function useUpdateEmail() {
             status: 'tokenRequired',
           }
         } else {
-          await updateEmailAndRefreshSession(agent, email, token)
+          await updateEmailAndRefreshSession(
+            agent,
+            partialRefreshSession,
+            email,
+            token,
+          )
           return {
             status: 'success',
           }

src/state/session/index.tsx

@@ -284,6 +284,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       type: 'partial-refresh-session',
       accountDid: agent.session!.did,
       patch: {
+        email: data.email,
         emailConfirmed: data.emailConfirmed,
         emailAuthFactor: data.emailAuthFactor,
       },

src/state/session/reducer.ts

@@ -58,7 +58,10 @@ export type Action =
   | {
       type: 'partial-refresh-session'
       accountDid: string
-      patch: Pick<SessionAccount, 'emailConfirmed' | 'emailAuthFactor'>
+      patch: Pick<
+        SessionAccount,
+        'email' | 'emailConfirmed' | 'emailAuthFactor'
+      >
     }
 
 function createPublicAgentState(): AgentState {
@@ -239,6 +242,7 @@ let reducer = (state: State, action: Action): State => {
        * Only mutating values that are safe. Be very careful with this.
        */
       if (agent.session) {
+        agent.session.email = patch.email ?? agent.session.email
         agent.session.emailConfirmed =
           patch.emailConfirmed ?? agent.session.emailConfirmed
         agent.session.emailAuthFactor =
@@ -255,6 +259,7 @@ let reducer = (state: State, action: Action): State => {
           if (a.did === accountDid) {
             return {
               ...a,
+              email: patch.email ?? a.email,
               emailConfirmed: patch.emailConfirmed ?? a.emailConfirmed,
               emailAuthFactor: patch.emailAuthFactor ?? a.emailAuthFactor,
             }