[Session V2] Persist agent changes automatically

src/state/session/index.tsx

@@ -8,11 +8,7 @@ import {logEvent} from '#/lib/statsig/statsig'
 import {logger} from '#/logger'
 import {isWeb} from '#/platform/detection'
 import * as persisted from '#/state/persisted'
-import {
+import {SessionApiContext, SessionStateContext} from '#/state/session/types'
-  SessionAccount,
-  SessionApiContext,
-  SessionStateContext,
-} from '#/state/session/types'
 import {
   agentToSessionAccount,
   configureModerationForAccount,
@@ -62,12 +58,12 @@ const ApiContext = React.createContext<SessionApiContext>({
 })
 
 export function Provider({children}: React.PropsWithChildren<{}>) {
-  const isDirty = React.useRef(false)
+  const [state, setState] = React.useState(() => ({
-  const [currentAgent, setCurrentAgent] =
+    accounts: persisted.get('session').accounts,
-    React.useState<BskyAgent>(INITIAL_AGENT)
+    currentAgent: INITIAL_AGENT,
-  const [accounts, setAccounts] = React.useState<SessionAccount[]>(
+    needsPersist: false,
-    persisted.get('session').accounts,
+  }))
-  )
+  const {accounts, currentAgent} = state
   const [isInitialLoad, setIsInitialLoad] = React.useState(true)
   const [isSwitchingAccounts, setIsSwitchingAccounts] = React.useState(false)
   const currentAccountDid = React.useMemo(
@@ -79,21 +75,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     [accounts, currentAccountDid],
   )
 
-  const persistNextUpdate = React.useCallback(
+  const switchAccountAndPersist = React.useCallback((newAgent: BskyAgent) => {
-    () => (isDirty.current = true),
+    const account = agentToSessionAccount(newAgent)
-    [],
+    setState(prev => {
-  )
+      let nextAccounts = prev.accounts
-
+      if (account) {
-  const upsertAndPersistAccount = React.useCallback(
+        nextAccounts = [
-    (account: SessionAccount) => {
-      persistNextUpdate()
-      setAccounts(accounts => [
           account,
-        ...accounts.filter(a => a.did !== account.did),
+          ...prev.accounts.filter(a => a.did !== account.did),
-      ])
+        ]
-    },
+      }
-    [setAccounts, persistNextUpdate],
+      return {
-  )
+        accounts: nextAccounts,
+        currentAgent: newAgent,
+        needsPersist: true,
+      }
+    })
+  }, [])
 
   const clearCurrentAccount = React.useCallback(() => {
     logger.warn(`session: clear current account`)
@@ -101,12 +99,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     // immediate clear this so any pending requests don't use it
     currentAgent.setPersistSessionHandler(() => {})
 
-    persistNextUpdate()
-
     const newAgent = new BskyAgent({service: PUBLIC_BSKY_SERVICE})
     configureModerationForGuest()
-    setCurrentAgent(newAgent)
+    switchAccountAndPersist(newAgent)
-  }, [currentAgent, persistNextUpdate, setCurrentAgent])
+  }, [currentAgent, switchAccountAndPersist])
 
   React.useEffect(() => {
     /*
@@ -161,14 +157,12 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
        * The updated account object, derived from the updated session we just
        * received from this callback.
        */
-      const refreshedAccount = agentToSessionAccount(currentAgent)
+      if (currentAgent.session) {
-
-      if (refreshedAccount) {
         /*
          * If the session expired naturally, or it was otherwise successfully
          * created/updated, we want to update/persist the data.
          */
-        upsertAndPersistAccount(refreshedAccount)
+        switchAccountAndPersist(currentAgent)
       } else {
         /*
          * This should never happen based on current `AtpAgent` handling, but
@@ -183,7 +177,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         }, 100)
       }
     })
-  }, [currentAgent, clearCurrentAccount, upsertAndPersistAccount])
+  }, [currentAgent, clearCurrentAccount, switchAccountAndPersist])
 
   const createAccount = React.useCallback<SessionApiContext['createAccount']>(
     async ({
@@ -199,8 +193,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       track('Try Create Account')
       logEvent('account:create:begin', {})
 
-      const {agent, account, fetchingGates} = await createAgentAndCreateAccount(
+      const {agent, fetchingGates} = await createAgentAndCreateAccount({
-        {
         service,
         email,
         password,
@@ -208,25 +201,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         inviteCode,
         verificationPhone,
         verificationCode,
-        },
+      })
-      )
 
       await fetchingGates
-      setCurrentAgent(agent)
+      switchAccountAndPersist(agent)
-      upsertAndPersistAccount(account)
 
       logger.debug(`session: created account`, {}, logger.DebugContext.session)
       track('Create Account')
       logEvent('account:create:success', {})
     },
-    [upsertAndPersistAccount],
+    [switchAccountAndPersist],
   )
 
   const login = React.useCallback<SessionApiContext['login']>(
     async ({service, identifier, password, authFactorToken}, logContext) => {
       logger.debug(`session: login`, {}, logger.DebugContext.session)
 
-      const {agent, account, fetchingGates} = await createAgentAndLogin({
+      const {agent, fetchingGates} = await createAgentAndLogin({
         service,
         identifier,
         password,
@@ -234,14 +225,13 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       })
 
       await fetchingGates
-      setCurrentAgent(agent)
+      switchAccountAndPersist(agent)
-      upsertAndPersistAccount(account)
 
       logger.debug(`session: logged in`, {}, logger.DebugContext.session)
       track('Sign In', {resumedSession: false})
       logEvent('account:loggedIn', {logContext, withPassword: true})
     },
-    [upsertAndPersistAccount],
+    [switchAccountAndPersist],
   )
 
   const logout = React.useCallback<SessionApiContext['logout']>(
@@ -249,18 +239,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       logger.debug(`session: logout`)
 
       clearCurrentAccount()
-      persistNextUpdate()
+      setState(prev => ({
-      setAccounts(accounts =>
+        accounts: prev.accounts.map(a => ({
-        accounts.map(a => ({
           ...a,
           accessJwt: undefined,
           refreshJwt: undefined,
         })),
-      )
+        currentAgent: prev.currentAgent,
+        needsPersist: true,
+      }))
 
       logEvent('account:loggedOut', {logContext})
     },
-    [clearCurrentAccount, persistNextUpdate, setAccounts],
+    [clearCurrentAccount],
   )
 
   const initSession = React.useCallback<SessionApiContext['initSession']>(
@@ -302,8 +293,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           logger.DebugContext.session,
         )
         await networkRetry(1, () => newAgent.resumeSession(prevSession))
-        setCurrentAgent(newAgent)
+        switchAccountAndPersist(newAgent)
-        upsertAndPersistAccount(agentToSessionAccount(newAgent)!)
       } else {
         /*
          * If the session is not expired, assume we can reuse it.
@@ -314,11 +304,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           logger.DebugContext.session,
         )
         newAgent.session = prevSession
-        setCurrentAgent(newAgent)
+        switchAccountAndPersist(newAgent)
-        upsertAndPersistAccount(account)
       }
     },
-    [upsertAndPersistAccount],
+    [switchAccountAndPersist],
   )
 
   const resumeSession = React.useCallback<SessionApiContext['resumeSession']>(
@@ -338,10 +327,13 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
   const removeAccount = React.useCallback<SessionApiContext['removeAccount']>(
     account => {
-      persistNextUpdate()
+      setState(prev => ({
-      setAccounts(accounts => accounts.filter(a => a.did !== account.did))
+        accounts: prev.accounts.filter(a => a.did !== account.did),
+        currentAgent: prev.currentAgent,
+        needsPersist: true,
+      }))
     },
-    [setAccounts, persistNextUpdate],
+    [],
   )
 
   const refreshSession = React.useCallback<
@@ -358,16 +350,8 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     await newAgent.resumeSession(sessionAccountToAgentSession(selectedAccount)!)
     const refreshedAccount = agentToSessionAccount(newAgent)
     await configureModerationForAccount(newAgent, refreshedAccount!)
-    persistNextUpdate()
+    switchAccountAndPersist(newAgent)
-    upsertAndPersistAccount(refreshedAccount!)
+  }, [currentAccountDid, currentAgent, switchAccountAndPersist])
-    setCurrentAgent(newAgent)
-  }, [
-    currentAccountDid,
-    currentAgent,
-    setCurrentAgent,
-    persistNextUpdate,
-    upsertAndPersistAccount,
-  ])
 
   const updateCurrentAccount = React.useCallback(async () => {
     await refreshSession()
@@ -391,14 +375,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   )
 
   React.useEffect(() => {
-    if (isDirty.current) {
+    if (state.needsPersist) {
-      isDirty.current = false
+      state.needsPersist = false
+
+      const currentAccountDid = state.currentAgent.session?.did
+      const currentAccount = state.accounts.find(
+        a => a.did === currentAccountDid,
+      )
       persisted.write('session', {
-        accounts,
+        accounts: state.accounts,
-        currentAccount,
+        currentAccount: currentAccount,
       })
     }
-  }, [accounts, currentAccount])
+  }, [state])
 
   React.useEffect(() => {
     return persisted.onUpdate(async () => {
@@ -414,7 +403,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
        * Accounts are already persisted on other side of broadcast, but we need
        * to update them in memory in this tab.
        */
-      setAccounts(persistedSession.accounts)
+      setState(prev => ({
+        accounts: persistedSession.accounts,
+        currentAgent: prev.currentAgent,
+        needsPersist: false, // We're syncing with another tab which already did that.
+      }))
 
       const selectedAccount = persistedSession.accounts.find(
         a => a.did === persistedSession.currentAccount?.did,
@@ -450,7 +443,11 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           const newAgent = currentAgent.clone()
           newAgent.session = sessionAccountToAgentSession(selectedAccount)
           await configureModerationForAccount(newAgent, selectedAccount)
-          setCurrentAgent(newAgent)
+          setState(prev => ({
+            accounts: prev.accounts,
+            currentAgent: newAgent,
+            needsPersist: false, // We're syncing with another tab which already did that.
+          }))
         }
       } else if (!selectedAccount && currentAccountDid) {
         logger.debug(
@@ -468,14 +465,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         clearCurrentAccount()
       }
     })
-  }, [
+  }, [currentAccountDid, clearCurrentAccount, initSession, currentAgent])
-    currentAccountDid,
-    setAccounts,
-    clearCurrentAccount,
-    initSession,
-    currentAgent,
-    setCurrentAgent,
-  ])
 
   const stateContext = React.useMemo(
     () => ({

src/state/session/util/index.ts

@@ -136,7 +136,6 @@ export async function createAgentAndLogin({
 
   return {
     agent,
-    account,
     fetchingGates,
   }
 }
@@ -185,7 +184,6 @@ export async function createAgentAndCreateAccount({
 
   return {
     agent,
-    account,
     fetchingGates,
   }
 }