Replace `resumeSession` with `getSession` in the email check (#8670)

* replace resumeSession with getSession

* copy lil type tweak from the other PR

* Add partialRefreshSession to session API context, use session state to infer state further down tree

* Review

---------

Co-authored-by: Eric Bailey <git@esb.lol>

7 files changed, 98 insertions, 65 deletions

diff --git a/src/components/dialogs/EmailDialog/data/useAccountEmailState.ts b/src/components/dialogs/EmailDialog/data/useAccountEmailState.ts
index 377411107..f25369f8d 100644
--- a/src/components/dialogs/EmailDialog/data/useAccountEmailState.ts
+++ b/src/components/dialogs/EmailDialog/data/useAccountEmailState.ts
@@ -1,7 +1,7 @@
-import {useCallback, useEffect, useState} from 'react'
-import {useQuery, useQueryClient} from '@tanstack/react-query'
+import {useEffect, useMemo, useState} from 'react'
+import {useQuery} from '@tanstack/react-query'
 
-import {useAgent} from '#/state/session'
+import {useAgent, useSessionApi} from '#/state/session'
 import {emitEmailVerified} from '#/components/dialogs/EmailDialog/events'
 
 export type AccountEmailState = {
@@ -11,57 +11,36 @@ export type AccountEmailState = {
 
 export const accountEmailStateQueryKey = ['accountEmailState'] as const
 
-export function useInvalidateAccountEmailState() {
-  const qc = useQueryClient()
-
-  return useCallback(() => {
-    return qc.invalidateQueries({
-      queryKey: accountEmailStateQueryKey,
-    })
-  }, [qc])
-}
-
-export function useUpdateAccountEmailStateQueryCache() {
-  const qc = useQueryClient()
-
-  return useCallback(
-    (data: AccountEmailState) => {
-      return qc.setQueriesData(
-        {
-          queryKey: accountEmailStateQueryKey,
-        },
-        data,
-      )
-    },
-    [qc],
-  )
-}
-
 export function useAccountEmailState() {
   const agent = useAgent()
+  const {partialRefreshSession} = useSessionApi()
   const [prevIsEmailVerified, setPrevEmailIsVerified] = useState(
     !!agent.session?.emailConfirmed,
   )
-  const fallbackData: AccountEmailState = {
-    isEmailVerified: !!agent.session?.emailConfirmed,
-    email2FAEnabled: !!agent.session?.emailAuthFactor,
-  }
-  const query = useQuery<AccountEmailState>({
+  const state: AccountEmailState = useMemo(
+    () => ({
+      isEmailVerified: !!agent.session?.emailConfirmed,
+      email2FAEnabled: !!agent.session?.emailAuthFactor,
+    }),
+    [agent.session],
+  )
+
+  /**
+   * Only here to refetch on focus, when necessary
+   */
+  useQuery({
     enabled: !!agent.session,
-    refetchOnWindowFocus: true,
+    /**
+     * Only refetch if the email verification s incomplete.
+     */
+    refetchOnWindowFocus: !prevIsEmailVerified,
     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,
-      }
+      await partialRefreshSession()
+      return null
     },
   })
 
-  const state = query.data ?? fallbackData
-
   /*
    * This will emit `n` times for each instance of this hook. So the listeners
    * all use `once` to prevent multiple handlers firing.
diff --git a/src/components/dialogs/EmailDialog/data/useConfirmEmail.ts b/src/components/dialogs/EmailDialog/data/useConfirmEmail.ts
index 73f824fcc..475a8cbfb 100644
--- a/src/components/dialogs/EmailDialog/data/useConfirmEmail.ts
+++ b/src/components/dialogs/EmailDialog/data/useConfirmEmail.ts
@@ -1,13 +1,10 @@
 import {useMutation} from '@tanstack/react-query'
 
 import {useAgent, useSession} from '#/state/session'
-import {useUpdateAccountEmailStateQueryCache} from '#/components/dialogs/EmailDialog/data/useAccountEmailState'
 
 export function useConfirmEmail() {
   const agent = useAgent()
   const {currentAccount} = useSession()
-  const updateAccountEmailStateQueryCache =
-    useUpdateAccountEmailStateQueryCache()
 
   return useMutation({
     mutationFn: async ({token}: {token: string}) => {
@@ -19,11 +16,8 @@ export function useConfirmEmail() {
         email: currentAccount.email,
         token: token.trim(),
       })
-      const {data} = await agent.resumeSession(agent.session!)
-      updateAccountEmailStateQueryCache({
-        isEmailVerified: !!data.emailConfirmed,
-        email2FAEnabled: !!data.emailAuthFactor,
-      })
+      // will update session state at root of app
+      await agent.resumeSession(agent.session!)
     },
   })
 }
diff --git a/src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts b/src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts
index 39f5fd2d9..358bf8654 100644
--- a/src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts
+++ b/src/components/dialogs/EmailDialog/data/useManageEmail2FA.ts
@@ -1,13 +1,10 @@
 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 ({
@@ -25,11 +22,8 @@ export function useManageEmail2FA() {
         emailAuthFactor: enabled,
         token,
       })
-      const {data} = await agent.resumeSession(agent.session!)
-      updateAccountEmailStateQueryCache({
-        isEmailVerified: !!data.emailConfirmed,
-        email2FAEnabled: !!data.emailAuthFactor,
-      })
+      // will update session state at root of app
+      await agent.resumeSession(agent.session!)
     },
   })
 }
diff --git a/src/lib/api/feed/home.ts b/src/lib/api/feed/home.ts
index e6bc45bea..7a0d72d91 100644
--- a/src/lib/api/feed/home.ts
+++ b/src/lib/api/feed/home.ts
@@ -1,9 +1,9 @@
-import {AppBskyFeedDefs, BskyAgent} from '@atproto/api'
+import {type AppBskyFeedDefs, type BskyAgent} from '@atproto/api'
 
 import {PROD_DEFAULT_FEED} from '#/lib/constants'
 import {CustomFeedAPI} from './custom'
 import {FollowingFeedAPI} from './following'
-import {FeedAPI, FeedAPIResponse} from './types'
+import {type FeedAPI, type FeedAPIResponse} from './types'
 
 // HACK
 // the feed API does not include any facilities for passing down
@@ -93,7 +93,7 @@ export class HomeFeedAPI implements FeedAPI {
       }
     }
 
-    if (this.usingDiscover) {
+    if (this.usingDiscover && !__DEV__) {
       const res = await this.discover.fetch({cursor, limit})
       returnCursor = res.cursor
       posts = posts.concat(res.feed)
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 45384c4f5..8223a7b3a 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -40,6 +40,7 @@ const ApiContext = React.createContext<SessionApiContext>({
   logoutEveryAccount: async () => {},
   resumeSession: async () => {},
   removeAccount: () => {},
+  partialRefreshSession: async () => {},
 })
 
 export function Provider({children}: React.PropsWithChildren<{}>) {
@@ -119,7 +120,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   )
 
   const logoutCurrentAccount = React.useCallback<
-    SessionApiContext['logoutEveryAccount']
+    SessionApiContext['logoutCurrentAccount']
   >(
     logContext => {
       addSessionDebugLog({type: 'method:start', method: 'logout'})
@@ -182,6 +183,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     [onAgentSessionChange, cancelPendingTask],
   )
 
+  const partialRefreshSession = React.useCallback<
+    SessionApiContext['partialRefreshSession']
+  >(async () => {
+    const agent = state.currentAgentState.agent as BskyAppAgent
+    const signal = cancelPendingTask()
+    const {data} = await agent.com.atproto.server.getSession()
+    if (signal.aborted) return
+    dispatch({
+      type: 'partial-refresh-session',
+      accountDid: agent.session!.did,
+      patch: {
+        emailConfirmed: data.emailConfirmed,
+        emailAuthFactor: data.emailAuthFactor,
+      },
+    })
+  }, [state, cancelPendingTask])
+
   const removeAccount = React.useCallback<SessionApiContext['removeAccount']>(
     account => {
       addSessionDebugLog({
@@ -262,6 +280,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       logoutEveryAccount,
       resumeSession,
       removeAccount,
+      partialRefreshSession,
     }),
     [
       createAccount,
@@ -270,6 +289,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       logoutEveryAccount,
       resumeSession,
       removeAccount,
+      partialRefreshSession,
     ],
   )
 
diff --git a/src/state/session/reducer.ts b/src/state/session/reducer.ts
index 22ba47162..f6452a391 100644
--- a/src/state/session/reducer.ts
+++ b/src/state/session/reducer.ts
@@ -1,8 +1,8 @@
-import {AtpSessionEvent} from '@atproto/api'
+import {type AtpSessionEvent, type BskyAgent} from '@atproto/api'
 
 import {createPublicAgent} from './agent'
 import {wrapSessionReducerForLogging} from './logging'
-import {SessionAccount} from './types'
+import {type SessionAccount} from './types'
 
 // A hack so that the reducer can't read anything from the agent.
 // From the reducer's point of view, it should be a completely opaque object.
@@ -52,6 +52,11 @@ export type Action =
       syncedAccounts: SessionAccount[]
       syncedCurrentDid: string | undefined
     }
+  | {
+      type: 'partial-refresh-session'
+      accountDid: string
+      patch: Pick<SessionAccount, 'emailConfirmed' | 'emailAuthFactor'>
+    }
 
 function createPublicAgentState(): AgentState {
   return {
@@ -180,6 +185,39 @@ let reducer = (state: State, action: Action): State => {
         needsPersist: false, // Synced from another tab. Don't persist to avoid cycles.
       }
     }
+    case 'partial-refresh-session': {
+      const {accountDid, patch} = action
+      const agent = state.currentAgentState.agent as BskyAgent
+
+      /*
+       * Only mutating values that are safe. Be very careful with this.
+       */
+      if (agent.session) {
+        agent.session.emailConfirmed =
+          patch.emailConfirmed ?? agent.session.emailConfirmed
+        agent.session.emailAuthFactor =
+          patch.emailAuthFactor ?? agent.session.emailAuthFactor
+      }
+
+      return {
+        ...state,
+        currentAgentState: {
+          ...state.currentAgentState,
+          agent,
+        },
+        accounts: state.accounts.map(a => {
+          if (a.did === accountDid) {
+            return {
+              ...a,
+              emailConfirmed: patch.emailConfirmed ?? a.emailConfirmed,
+              emailAuthFactor: patch.emailAuthFactor ?? a.emailAuthFactor,
+            }
+          }
+          return a
+        }),
+        needsPersist: true,
+      }
+    }
   }
 }
 reducer = wrapSessionReducerForLogging(reducer)
diff --git a/src/state/session/types.ts b/src/state/session/types.ts
index aa8b9a99e..4621b4f04 100644
--- a/src/state/session/types.ts
+++ b/src/state/session/types.ts
@@ -40,4 +40,12 @@ export type SessionApiContext = {
   ) => void
   resumeSession: (account: SessionAccount) => Promise<void>
   removeAccount: (account: SessionAccount) => void
+  /**
+   * Calls `getSession` and updates select fields on the current account and
+   * `BskyAgent`. This is an alternative to `resumeSession`, which updates
+   * current account/agent using the `persistSessionHandler`, but is more load
+   * bearing. This patches in updates without causing any side effects via
+   * `persistSessionHandler`.
+   */
+  partialRefreshSession: () => Promise<void>
 }