[Session] Track agent in state (#3833)

* Expand currentAccountDid into currentAgentState

* Inline all callsites of upsertAccount

* Inline all internal callsites of clearCurrentAccount

* Add agent to currentAgentState

* Collapse extra setStates

* Add TODOs

src/state/session/index.tsx

@@ -56,31 +56,26 @@ function __getAgent() {
   return __globalAgent
 }
 
+type AgentState = {
+  readonly agent: BskyAgent
+  readonly did: string | undefined
+}
+
 type State = {
   accounts: SessionStateContext['accounts']
-  currentAccountDid: string | undefined
+  currentAgentState: AgentState
   needsPersist: boolean
 }
 
 export function Provider({children}: React.PropsWithChildren<{}>) {
-  const [state, setState] = React.useState<State>({
+  const [state, setState] = React.useState<State>(() => ({
     accounts: persisted.get('session').accounts,
-    currentAccountDid: undefined, // assume logged out to start
+    currentAgentState: {
-    needsPersist: false,
+      agent: PUBLIC_BSKY_AGENT,
-  })
+      did: undefined, // assume logged out to start
-
-  const upsertAccount = React.useCallback(
-    (account: SessionAccount, expired = false) => {
-      setState(s => {
-        return {
-          accounts: [account, ...s.accounts.filter(a => a.did !== account.did)],
-          currentAccountDid: expired ? undefined : account.did,
-          needsPersist: true,
-        }
-      })
     },
-    [setState],
+    needsPersist: false,
-  )
+  }))
 
   const clearCurrentAccount = React.useCallback(() => {
     logger.warn(`session: clear current account`)
@@ -88,7 +83,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     configureModerationForGuest()
     setState(s => ({
       accounts: s.accounts,
-      currentAccountDid: undefined,
+      currentAgentState: {
+        agent: PUBLIC_BSKY_AGENT,
+        did: undefined,
+      },
       needsPersist: true,
     }))
   }, [setState])
@@ -106,7 +104,17 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         logger.warn(
           `session: persistSessionHandler received network-error event`,
         )
-        clearCurrentAccount()
+        logger.warn(`session: clear current account`)
+        __globalAgent = PUBLIC_BSKY_AGENT
+        configureModerationForGuest()
+        setState(s => ({
+          accounts: s.accounts,
+          currentAgentState: {
+            agent: PUBLIC_BSKY_AGENT,
+            did: undefined,
+          },
+          needsPersist: true,
+        }))
         return
       }
 
@@ -149,9 +157,23 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
        * to persist this data and wipe their tokens, effectively logging them
        * out.
        */
-      upsertAccount(refreshedAccount, expired)
+      setState(s => {
+        return {
+          accounts: [
+            refreshedAccount,
+            ...s.accounts.filter(a => a.did !== refreshedAccount.did),
+          ],
+          currentAgentState: expired
+            ? {
+                agent: PUBLIC_BSKY_AGENT,
+                did: undefined,
+              }
+            : s.currentAgentState,
+          needsPersist: true,
+        }
+      })
     },
-    [clearCurrentAccount, upsertAccount],
+    [],
   )
 
   const createAccount = React.useCallback<SessionApiContext['createAccount']>(
@@ -185,13 +207,22 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
       __globalAgent = agent
       await fetchingGates
-      upsertAccount(account)
+      setState(s => {
+        return {
+          accounts: [account, ...s.accounts.filter(a => a.did !== account.did)],
+          currentAgentState: {
+            did: account.did,
+            agent: agent,
+          },
+          needsPersist: true,
+        }
+      })
 
       logger.debug(`session: created account`, {}, logger.DebugContext.session)
       track('Create Account')
       logEvent('account:create:success', {})
     },
-    [upsertAccount, onAgentSessionChange],
+    [onAgentSessionChange],
   )
 
   const login = React.useCallback<SessionApiContext['login']>(
@@ -212,20 +243,31 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       // @ts-ignore
       if (IS_DEV && isWeb) window.agent = agent
       await fetchingGates
-      upsertAccount(account)
+      setState(s => {
+        return {
+          accounts: [account, ...s.accounts.filter(a => a.did !== account.did)],
+          currentAgentState: {
+            did: account.did,
+            agent: agent,
+          },
+          needsPersist: true,
+        }
+      })
 
       logger.debug(`session: logged in`, {}, logger.DebugContext.session)
 
       track('Sign In', {resumedSession: false})
       logEvent('account:loggedIn', {logContext, withPassword: true})
     },
-    [upsertAccount, onAgentSessionChange],
+    [onAgentSessionChange],
   )
 
   const logout = React.useCallback<SessionApiContext['logout']>(
     async logContext => {
       logger.debug(`session: logout`)
-      clearCurrentAccount()
+      logger.warn(`session: clear current account`)
+      __globalAgent = PUBLIC_BSKY_AGENT
+      configureModerationForGuest()
       setState(s => {
         return {
           accounts: s.accounts.map(a => ({
@@ -233,13 +275,16 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
             refreshJwt: undefined,
             accessJwt: undefined,
           })),
-          currentAccountDid: s.currentAccountDid,
+          currentAgentState: {
+            did: undefined,
+            agent: PUBLIC_BSKY_AGENT,
+          },
           needsPersist: true,
         }
       })
       logEvent('account:loggedOut', {logContext})
     },
-    [clearCurrentAccount, setState],
+    [setState],
   )
 
   const initSession = React.useCallback<SessionApiContext['initSession']>(
@@ -279,7 +324,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           const freshAccount = await resumeSessionWithFreshAccount()
           __globalAgent = agent
           await fetchingGates
-          upsertAccount(freshAccount)
+          setState(s => {
+            return {
+              accounts: [
+                freshAccount,
+                ...s.accounts.filter(a => a.did !== freshAccount.did),
+              ],
+              currentAgentState: {
+                did: freshAccount.did,
+                agent: agent,
+              },
+              needsPersist: true,
+            }
+          })
         } catch (e) {
           /*
            * Note: `agent.persistSession` is also called when this fails, and
@@ -290,7 +347,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           })
 
           __globalAgent = PUBLIC_BSKY_AGENT
-          // TODO: Should this update currentAccountDid?
+          // TODO: This needs a setState.
         }
       } else {
         logger.debug(`session: attempting to reuse previous session`)
@@ -299,7 +356,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
         __globalAgent = agent
         await fetchingGates
-        upsertAccount(account)
+        setState(s => {
+          return {
+            accounts: [
+              account,
+              ...s.accounts.filter(a => a.did !== account.did),
+            ],
+            currentAgentState: {
+              did: account.did,
+              agent: agent,
+            },
+            needsPersist: true,
+          }
+        })
 
         if (accountOrSessionDeactivated) {
           // don't attempt to resume
@@ -315,7 +384,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
               logger.info(
                 `session: reuse of previous session returned a fresh account, upserting`,
               )
-              upsertAccount(freshAccount)
+              setState(s => {
+                return {
+                  accounts: [
+                    freshAccount,
+                    ...s.accounts.filter(a => a.did !== freshAccount.did),
+                  ],
+                  currentAgentState: {
+                    did: freshAccount.did,
+                    agent: agent,
+                  },
+                  needsPersist: true,
+                }
+              })
             }
           })
           .catch(e => {
@@ -328,7 +409,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
             })
 
             __globalAgent = PUBLIC_BSKY_AGENT
-            // TODO: Should this update currentAccountDid?
+            // TODO: This needs a setState.
           })
       }
 
@@ -347,7 +428,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         return sessionAccount
       }
     },
-    [upsertAccount, onAgentSessionChange],
+    [onAgentSessionChange],
   )
 
   const removeAccount = React.useCallback<SessionApiContext['removeAccount']>(
@@ -355,7 +436,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       setState(s => {
         return {
           accounts: s.accounts.filter(a => a.did !== account.did),
-          currentAccountDid: s.currentAccountDid,
+          currentAgentState: s.currentAgentState,
           needsPersist: true,
         }
       })
@@ -369,7 +450,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     account => {
       setState(s => {
         const currentAccount = s.accounts.find(
-          a => a.did === s.currentAccountDid,
+          a => a.did === s.currentAgentState.did,
         )
         // ignore, should never happen
         if (!currentAccount) return s
@@ -393,7 +474,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
             updatedAccount,
             ...s.accounts.filter(a => a.did !== currentAccount.did),
           ],
-          currentAccountDid: s.currentAccountDid,
+          currentAgentState: s.currentAgentState,
           needsPersist: true,
         }
       })
@@ -407,7 +488,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       persisted.write('session', {
         accounts: state.accounts,
         currentAccount: state.accounts.find(
-          a => a.did === state.currentAccountDid,
+          a => a.did === state.currentAgentState.did,
         ),
       })
     }
@@ -424,10 +505,10 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       )
 
       if (selectedAccount && selectedAccount.refreshJwt) {
-        if (selectedAccount.did !== state.currentAccountDid) {
+        if (selectedAccount.did !== state.currentAgentState.did) {
           logger.debug(`session: persisted onUpdate, switching accounts`, {
             from: {
-              did: state.currentAccountDid,
+              did: state.currentAgentState.did,
             },
             to: {
               did: selectedAccount.did,
@@ -445,8 +526,9 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
            */
           // @ts-ignore we checked for `refreshJwt` above
           __globalAgent.session = selectedAccount
+          // TODO: This needs a setState.
         }
-      } else if (!selectedAccount && state.currentAccountDid) {
+      } else if (!selectedAccount && state.currentAgentState.did) {
         logger.debug(
           `session: persisted onUpdate, logging out`,
           {},
@@ -459,24 +541,34 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
          * handled by `persistSession` (which nukes this accounts tokens only),
          * or by a `logout` call  which nukes all accounts tokens)
          */
-        clearCurrentAccount()
+        logger.warn(`session: clear current account`)
+        __globalAgent = PUBLIC_BSKY_AGENT
+        configureModerationForGuest()
+        setState(s => ({
+          accounts: s.accounts,
+          currentAgentState: {
+            did: undefined,
+            agent: PUBLIC_BSKY_AGENT,
+          },
+          needsPersist: true, // TODO: This seems bad in this codepath. Existing behavior.
+        }))
       }
 
-      setState(() => ({
+      setState(s => ({
         accounts: persistedSession.accounts,
-        currentAccountDid: selectedAccount?.did,
+        currentAgentState: s.currentAgentState,
         needsPersist: false, // Synced from another tab. Don't persist to avoid cycles.
       }))
     })
-  }, [state, setState, clearCurrentAccount, initSession])
+  }, [state, setState, initSession])
 
   const stateContext = React.useMemo(
     () => ({
       accounts: state.accounts,
       currentAccount: state.accounts.find(
-        a => a.did === state.currentAccountDid,
+        a => a.did === state.currentAgentState.did,
       ),
-      hasSession: !!state.currentAccountDid,
+      hasSession: !!state.currentAgentState.did,
     }),
     [state],
   )