[Session] Code cleanup (#3854)

* Split utils into files

* Move reducer to another file

* Write types explicitly

* Remove unnnecessary check

* Move things around a bit

* Move more stuff into agent factories

* Move more stuff into agent

* Fix gates await

* Clarify comments

* Enforce more via types

* Nit

* initSession -> resumeSession

* Protect against races

* Make agent opaque to reducer

* Check using plain condition

src/App.native.tsx

@@ -57,7 +57,7 @@ SplashScreen.preventAutoHideAsync()
 function InnerApp() {
   const [isReady, setIsReady] = React.useState(false)
   const {currentAccount} = useSession()
-  const {initSession} = useSessionApi()
+  const {resumeSession} = useSessionApi()
   const theme = useColorModeTheme()
   const {_} = useLingui()
 
@@ -65,20 +65,20 @@ function InnerApp() {
 
   // init
   useEffect(() => {
-    async function resumeSession(account?: SessionAccount) {
+    async function onLaunch(account?: SessionAccount) {
       try {
         if (account) {
-          await initSession(account)
+          await resumeSession(account)
         }
       } catch (e) {
-        logger.error(`session: resumeSession failed`, {message: e})
+        logger.error(`session: resume failed`, {message: e})
       } finally {
         setIsReady(true)
       }
     }
     const account = readLastActiveAccount()
-    resumeSession(account)
+    onLaunch(account)
-  }, [initSession])
+  }, [resumeSession])
 
   useEffect(() => {
     return listenSessionDropped(() => {

src/App.web.tsx

@@ -45,17 +45,17 @@ import {listenSessionDropped} from './state/events'
 function InnerApp() {
   const [isReady, setIsReady] = React.useState(false)
   const {currentAccount} = useSession()
-  const {initSession} = useSessionApi()
+  const {resumeSession} = useSessionApi()
   const theme = useColorModeTheme()
   const {_} = useLingui()
   useIntentHandler()
 
   // init
   useEffect(() => {
-    async function resumeSession(account?: SessionAccount) {
+    async function onLaunch(account?: SessionAccount) {
       try {
         if (account) {
-          await initSession(account)
+          await resumeSession(account)
         }
       } catch (e) {
         logger.error(`session: resumeSession failed`, {message: e})
@@ -64,8 +64,8 @@ function InnerApp() {
       }
     }
     const account = readLastActiveAccount()
-    resumeSession(account)
+    onLaunch(account)
-  }, [initSession])
+  }, [resumeSession])
 
   useEffect(() => {
     return listenSessionDropped(() => {

src/lib/hooks/useAccountSwitcher.ts

@@ -15,7 +15,7 @@ export function useAccountSwitcher() {
   const [pendingDid, setPendingDid] = useState<string | null>(null)
   const {_} = useLingui()
   const {track} = useAnalytics()
-  const {initSession} = useSessionApi()
+  const {resumeSession} = useSessionApi()
   const {requestSwitchToAccount} = useLoggedOutViewControls()
 
   const onPressSwitchAccount = useCallback(
@@ -39,7 +39,7 @@ export function useAccountSwitcher() {
             // So we change the URL ourselves. The navigator will pick it up on remount.
             history.pushState(null, '', '/')
           }
-          await initSession(account)
+          await resumeSession(account)
           logEvent('account:loggedIn', {logContext, withPassword: false})
           Toast.show(_(msg`Signed in as @${account.handle}`))
         } else {
@@ -57,7 +57,7 @@ export function useAccountSwitcher() {
         setPendingDid(null)
       }
     },
-    [_, track, initSession, requestSwitchToAccount, pendingDid],
+    [_, track, resumeSession, requestSwitchToAccount, pendingDid],
   )
 
   return {onPressSwitchAccount, pendingDid}

src/screens/Login/ChooseAccountForm.tsx

@@ -26,7 +26,7 @@ export const ChooseAccountForm = ({
   const {track, screen} = useAnalytics()
   const {_} = useLingui()
   const {currentAccount} = useSession()
-  const {initSession} = useSessionApi()
+  const {resumeSession} = useSessionApi()
   const {setShowLoggedOut} = useLoggedOutViewControls()
 
   React.useEffect(() => {
@@ -51,7 +51,7 @@ export const ChooseAccountForm = ({
       }
       try {
         setPendingDid(account.did)
-        await initSession(account)
+        await resumeSession(account)
         logEvent('account:loggedIn', {
           logContext: 'ChooseAccountForm',
           withPassword: false,
@@ -71,7 +71,7 @@ export const ChooseAccountForm = ({
     [
       currentAccount,
       track,
-      initSession,
+      resumeSession,
       pendingDid,
       onSelectAccount,
       setShowLoggedOut,

src/state/session/agent.ts

@@ -0,0 +1,190 @@
+import {BskyAgent} from '@atproto/api'
+import {AtpSessionEvent} from '@atproto-labs/api'
+
+import {networkRetry} from '#/lib/async/retry'
+import {PUBLIC_BSKY_SERVICE} from '#/lib/constants'
+import {tryFetchGates} from '#/lib/statsig/statsig'
+import {
+  configureModerationForAccount,
+  configureModerationForGuest,
+} from './moderation'
+import {SessionAccount} from './types'
+import {isSessionDeactivated, isSessionExpired} from './util'
+import {IS_PROD_SERVICE} from '#/lib/constants'
+import {DEFAULT_PROD_FEEDS} from '../queries/preferences'
+
+export function createPublicAgent() {
+  configureModerationForGuest() // Side effect but only relevant for tests
+  return new BskyAgent({service: PUBLIC_BSKY_SERVICE})
+}
+
+export async function createAgentAndResume(
+  storedAccount: SessionAccount,
+  onSessionChange: (
+    agent: BskyAgent,
+    did: string,
+    event: AtpSessionEvent,
+  ) => void,
+) {
+  const agent = new BskyAgent({service: storedAccount.service})
+  if (storedAccount.pdsUrl) {
+    agent.pdsUrl = agent.api.xrpc.uri = new URL(storedAccount.pdsUrl)
+  }
+  const gates = tryFetchGates(storedAccount.did, 'prefer-low-latency')
+  const moderation = configureModerationForAccount(agent, storedAccount)
+  const prevSession = {
+    accessJwt: storedAccount.accessJwt ?? '',
+    refreshJwt: storedAccount.refreshJwt ?? '',
+    did: storedAccount.did,
+    handle: storedAccount.handle,
+  }
+  if (isSessionExpired(storedAccount)) {
+    await networkRetry(1, () => agent.resumeSession(prevSession))
+  } else {
+    agent.session = prevSession
+    if (!storedAccount.deactivated) {
+      // Intentionally not awaited to unblock the UI:
+      networkRetry(1, () => agent.resumeSession(prevSession))
+    }
+  }
+
+  return prepareAgent(agent, gates, moderation, onSessionChange)
+}
+
+export async function createAgentAndLogin(
+  {
+    service,
+    identifier,
+    password,
+    authFactorToken,
+  }: {
+    service: string
+    identifier: string
+    password: string
+    authFactorToken?: string
+  },
+  onSessionChange: (
+    agent: BskyAgent,
+    did: string,
+    event: AtpSessionEvent,
+  ) => void,
+) {
+  const agent = new BskyAgent({service})
+  await agent.login({identifier, password, authFactorToken})
+
+  const account = agentToSessionAccountOrThrow(agent)
+  const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
+  const moderation = configureModerationForAccount(agent, account)
+  return prepareAgent(agent, moderation, gates, onSessionChange)
+}
+
+export async function createAgentAndCreateAccount(
+  {
+    service,
+    email,
+    password,
+    handle,
+    birthDate,
+    inviteCode,
+    verificationPhone,
+    verificationCode,
+  }: {
+    service: string
+    email: string
+    password: string
+    handle: string
+    birthDate: Date
+    inviteCode?: string
+    verificationPhone?: string
+    verificationCode?: string
+  },
+  onSessionChange: (
+    agent: BskyAgent,
+    did: string,
+    event: AtpSessionEvent,
+  ) => void,
+) {
+  const agent = new BskyAgent({service})
+  await agent.createAccount({
+    email,
+    password,
+    handle,
+    inviteCode,
+    verificationPhone,
+    verificationCode,
+  })
+  const account = agentToSessionAccountOrThrow(agent)
+  const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
+  const moderation = configureModerationForAccount(agent, account)
+  if (!account.deactivated) {
+    /*dont await*/ agent.upsertProfile(_existing => {
+      return {
+        displayName: '',
+        // HACKFIX
+        // creating a bunch of identical profile objects is breaking the relay
+        // tossing this unspecced field onto it to reduce the size of the problem
+        // -prf
+        createdAt: new Date().toISOString(),
+      }
+    })
+  }
+
+  // Not awaited so that we can still get into onboarding.
+  // This is OK because we won't let you toggle adult stuff until you set the date.
+  agent.setPersonalDetails({birthDate: birthDate.toISOString()})
+  if (IS_PROD_SERVICE(service)) {
+    agent.setSavedFeeds(DEFAULT_PROD_FEEDS.saved, DEFAULT_PROD_FEEDS.pinned)
+  }
+
+  return prepareAgent(agent, gates, moderation, onSessionChange)
+}
+
+async function prepareAgent(
+  agent: BskyAgent,
+  // Not awaited in the calling code so we can delay blocking on them.
+  gates: Promise<void>,
+  moderation: Promise<void>,
+  onSessionChange: (
+    agent: BskyAgent,
+    did: string,
+    event: AtpSessionEvent,
+  ) => void,
+) {
+  // There's nothing else left to do, so block on them here.
+  await Promise.all([gates, moderation])
+
+  // Now the agent is ready.
+  const account = agentToSessionAccountOrThrow(agent)
+  agent.setPersistSessionHandler(event => {
+    onSessionChange(agent, account.did, event)
+  })
+  return {agent, account}
+}
+
+export function agentToSessionAccountOrThrow(agent: BskyAgent): SessionAccount {
+  const account = agentToSessionAccount(agent)
+  if (!account) {
+    throw Error('Expected an active session')
+  }
+  return account
+}
+
+export function agentToSessionAccount(
+  agent: BskyAgent,
+): SessionAccount | undefined {
+  if (!agent.session) {
+    return undefined
+  }
+  return {
+    service: agent.service.toString(),
+    did: agent.session.did,
+    handle: agent.session.handle,
+    email: agent.session.email,
+    emailConfirmed: agent.session.emailConfirmed || false,
+    emailAuthFactor: agent.session.emailAuthFactor || false,
+    refreshJwt: agent.session.refreshJwt,
+    accessJwt: agent.session.accessJwt,
+    deactivated: isSessionDeactivated(agent.session.accessJwt),
+    pdsUrl: agent.pdsUrl?.toString(),
+  }
+}

src/state/session/index.tsx

@@ -2,9 +2,7 @@ import React from 'react'
 import {AtpSessionEvent, BskyAgent} from '@atproto/api'
 
 import {track} from '#/lib/analytics/analytics'
-import {networkRetry} from '#/lib/async/retry'
+import {logEvent} from '#/lib/statsig/statsig'
-import {PUBLIC_BSKY_SERVICE} from '#/lib/constants'
-import {logEvent, tryFetchGates} from '#/lib/statsig/statsig'
 import {isWeb} from '#/platform/detection'
 import * as persisted from '#/state/persisted'
 import {useCloseAllActiveElements} from '#/state/util'
@@ -13,22 +11,15 @@ import {IS_DEV} from '#/env'
 import {emitSessionDropped} from '../events'
 import {
   agentToSessionAccount,
-  configureModerationForAccount,
-  configureModerationForGuest,
   createAgentAndCreateAccount,
   createAgentAndLogin,
-  isSessionDeactivated,
+  createAgentAndResume,
-  isSessionExpired,
+} from './agent'
-} from './util'
+import {getInitialState, reducer} from './reducer'
 
+export {isSessionDeactivated} from './util'
 export type {SessionAccount} from '#/state/session/types'
-import {
+import {SessionApiContext, SessionStateContext} from '#/state/session/types'
-  SessionAccount,
-  SessionApiContext,
-  SessionStateContext,
-} from '#/state/session/types'
-
-export {isSessionDeactivated}
 
 const StateContext = React.createContext<SessionStateContext>({
   accounts: [],
@@ -42,190 +33,16 @@ const ApiContext = React.createContext<SessionApiContext>({
   createAccount: async () => {},
   login: async () => {},
   logout: async () => {},
-  initSession: async () => {},
+  resumeSession: async () => {},
   removeAccount: () => {},
   updateCurrentAccount: () => {},
 })
 
-type AgentState = {
-  readonly agent: BskyAgent
-  readonly did: string | undefined
-}
-
-type State = {
-  accounts: SessionStateContext['accounts']
-  currentAgentState: AgentState
-  needsPersist: boolean
-}
-
-type Action =
-  | {
-      type: 'received-agent-event'
-      agent: BskyAgent
-      accountDid: string
-      refreshedAccount: SessionAccount | undefined
-      sessionEvent: AtpSessionEvent
-    }
-  | {
-      type: 'switched-to-account'
-      newAgent: BskyAgent
-      newAccount: SessionAccount
-    }
-  | {
-      type: 'updated-current-account'
-      updatedFields: Partial<
-        Pick<
-          SessionAccount,
-          'handle' | 'email' | 'emailConfirmed' | 'emailAuthFactor'
-        >
-      >
-    }
-  | {
-      type: 'removed-account'
-      accountDid: string
-    }
-  | {
-      type: 'logged-out'
-    }
-  | {
-      type: 'synced-accounts'
-      syncedAccounts: SessionAccount[]
-      syncedCurrentDid: string | undefined
-    }
-
-function createPublicAgentState() {
-  configureModerationForGuest() // Side effect but only relevant for tests
-  return {
-    agent: new BskyAgent({service: PUBLIC_BSKY_SERVICE}),
-    did: undefined,
-  }
-}
-
-function getInitialState(): State {
-  return {
-    accounts: persisted.get('session').accounts,
-    currentAgentState: createPublicAgentState(),
-    needsPersist: false,
-  }
-}
-
-function reducer(state: State, action: Action): State {
-  switch (action.type) {
-    case 'received-agent-event': {
-      const {agent, accountDid, refreshedAccount, sessionEvent} = action
-      if (agent !== state.currentAgentState.agent) {
-        // Only consider events from the active agent.
-        return state
-      }
-      if (sessionEvent === 'network-error') {
-        // Don't change stored accounts but kick to the choose account screen.
-        return {
-          accounts: state.accounts,
-          currentAgentState: createPublicAgentState(),
-          needsPersist: true,
-        }
-      }
-      const existingAccount = state.accounts.find(a => a.did === accountDid)
-      if (
-        !existingAccount ||
-        JSON.stringify(existingAccount) === JSON.stringify(refreshedAccount)
-      ) {
-        // Fast path without a state update.
-        return state
-      }
-      return {
-        accounts: state.accounts.map(a => {
-          if (a.did === accountDid) {
-            if (refreshedAccount) {
-              return refreshedAccount
-            } else {
-              return {
-                ...a,
-                // If we didn't receive a refreshed account, clear out the tokens.
-                accessJwt: undefined,
-                refreshJwt: undefined,
-              }
-            }
-          } else {
-            return a
-          }
-        }),
-        currentAgentState: refreshedAccount
-          ? state.currentAgentState
-          : createPublicAgentState(), // Log out if expired.
-        needsPersist: true,
-      }
-    }
-    case 'switched-to-account': {
-      const {newAccount, newAgent} = action
-      return {
-        accounts: [
-          newAccount,
-          ...state.accounts.filter(a => a.did !== newAccount.did),
-        ],
-        currentAgentState: {
-          did: newAccount.did,
-          agent: newAgent,
-        },
-        needsPersist: true,
-      }
-    }
-    case 'updated-current-account': {
-      const {updatedFields} = action
-      return {
-        accounts: state.accounts.map(a => {
-          if (a.did === state.currentAgentState.did) {
-            return {
-              ...a,
-              ...updatedFields,
-            }
-          } else {
-            return a
-          }
-        }),
-        currentAgentState: state.currentAgentState,
-        needsPersist: true,
-      }
-    }
-    case 'removed-account': {
-      const {accountDid} = action
-      return {
-        accounts: state.accounts.filter(a => a.did !== accountDid),
-        currentAgentState:
-          state.currentAgentState.did === accountDid
-            ? createPublicAgentState() // Log out if removing the current one.
-            : state.currentAgentState,
-        needsPersist: true,
-      }
-    }
-    case 'logged-out': {
-      return {
-        accounts: state.accounts.map(a => ({
-          ...a,
-          // Clear tokens for *every* account (this is a hard logout).
-          refreshJwt: undefined,
-          accessJwt: undefined,
-        })),
-        currentAgentState: createPublicAgentState(),
-        needsPersist: true,
-      }
-    }
-    case 'synced-accounts': {
-      const {syncedAccounts, syncedCurrentDid} = action
-      return {
-        accounts: syncedAccounts,
-        currentAgentState:
-          syncedCurrentDid === state.currentAgentState.did
-            ? state.currentAgentState
-            : createPublicAgentState(), // Log out if different user.
-        needsPersist: false, // Synced from another tab. Don't persist to avoid cycles.
-      }
-    }
-  }
-}
-
 export function Provider({children}: React.PropsWithChildren<{}>) {
-  const [state, dispatch] = React.useReducer(reducer, null, getInitialState)
+  const cancelPendingTask = useOneTaskAtATime()
+  const [state, dispatch] = React.useReducer(reducer, null, () =>
+    getInitialState(persisted.get('session').accounts),
+  )
 
   const onAgentSessionChange = React.useCallback(
     (agent: BskyAgent, accountDid: string, sessionEvent: AtpSessionEvent) => {
@@ -245,34 +62,18 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   )
 
   const createAccount = React.useCallback<SessionApiContext['createAccount']>(
-    async ({
+    async params => {
-      service,
+      const signal = cancelPendingTask()
-      email,
-      password,
-      handle,
-      birthDate,
-      inviteCode,
-      verificationPhone,
-      verificationCode,
-    }) => {
       track('Try Create Account')
       logEvent('account:create:begin', {})
-      const {agent, account, fetchingGates} = await createAgentAndCreateAccount(
+      const {agent, account} = await createAgentAndCreateAccount(
-        {
+        params,
-          service,
+        onAgentSessionChange,
-          email,
-          password,
-          handle,
-          birthDate,
-          inviteCode,
-          verificationPhone,
-          verificationCode,
-        },
       )
-      agent.setPersistSessionHandler(event => {
+
-        onAgentSessionChange(agent, account.did, event)
+      if (signal.aborted) {
-      })
+        return
-      await fetchingGates
+      }
       dispatch({
         type: 'switched-to-account',
         newAgent: agent,
@@ -281,21 +82,20 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       track('Create Account')
       logEvent('account:create:success', {})
     },
-    [onAgentSessionChange],
+    [onAgentSessionChange, cancelPendingTask],
   )
 
   const login = React.useCallback<SessionApiContext['login']>(
-    async ({service, identifier, password, authFactorToken}, logContext) => {
+    async (params, logContext) => {
-      const {agent, account, fetchingGates} = await createAgentAndLogin({
+      const signal = cancelPendingTask()
-        service,
+      const {agent, account} = await createAgentAndLogin(
-        identifier,
+        params,
-        password,
+        onAgentSessionChange,
-        authFactorToken,
+      )
-      })
+
-      agent.setPersistSessionHandler(event => {
+      if (signal.aborted) {
-        onAgentSessionChange(agent, account.did, event)
+        return
-      })
+      }
-      await fetchingGates
       dispatch({
         type: 'switched-to-account',
         newAgent: agent,
@@ -304,88 +104,49 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       track('Sign In', {resumedSession: false})
       logEvent('account:loggedIn', {logContext, withPassword: true})
     },
-    [onAgentSessionChange],
+    [onAgentSessionChange, cancelPendingTask],
   )
 
   const logout = React.useCallback<SessionApiContext['logout']>(
-    async logContext => {
+    logContext => {
+      cancelPendingTask()
       dispatch({
         type: 'logged-out',
       })
       logEvent('account:loggedOut', {logContext})
     },
-    [],
+    [cancelPendingTask],
   )
 
-  const initSession = React.useCallback<SessionApiContext['initSession']>(
+  const resumeSession = React.useCallback<SessionApiContext['resumeSession']>(
-    async account => {
+    async storedAccount => {
-      const fetchingGates = tryFetchGates(account.did, 'prefer-low-latency')
+      const signal = cancelPendingTask()
-      const agent = new BskyAgent({service: account.service})
+      const {agent, account} = await createAgentAndResume(
-      // restore the correct PDS URL if available
+        storedAccount,
-      if (account.pdsUrl) {
+        onAgentSessionChange,
-        agent.pdsUrl = agent.api.xrpc.uri = new URL(account.pdsUrl)
+      )
-      }
-      agent.setPersistSessionHandler(event => {
-        onAgentSessionChange(agent, account.did, event)
-      })
-      await configureModerationForAccount(agent, account)
 
-      const prevSession = {
+      if (signal.aborted) {
-        accessJwt: account.accessJwt ?? '',
+        return
-        refreshJwt: account.refreshJwt ?? '',
-        did: account.did,
-        handle: account.handle,
       }
-
-      if (isSessionExpired(account)) {
-        const freshAccount = await resumeSessionWithFreshAccount()
-        await fetchingGates
-        dispatch({
-          type: 'switched-to-account',
-          newAgent: agent,
-          newAccount: freshAccount,
-        })
-      } else {
-        agent.session = prevSession
-        await fetchingGates
       dispatch({
         type: 'switched-to-account',
         newAgent: agent,
         newAccount: account,
       })
-        if (isSessionDeactivated(account.accessJwt) || account.deactivated) {
-          // don't attempt to resume
-          // use will be taken to the deactivated screen
-          return
-        }
-        // Intentionally not awaited to unblock the UI:
-        resumeSessionWithFreshAccount()
-      }
-
-      async function resumeSessionWithFreshAccount(): Promise<SessionAccount> {
-        await networkRetry(1, () => agent.resumeSession(prevSession))
-        const sessionAccount = agentToSessionAccount(agent)
-        /*
-         * If `agent.resumeSession` fails above, it'll throw. This is just to
-         * make TypeScript happy.
-         */
-        if (!sessionAccount) {
-          throw new Error(`session: initSession failed to establish a session`)
-        }
-        return sessionAccount
-      }
     },
-    [onAgentSessionChange],
+    [onAgentSessionChange, cancelPendingTask],
   )
 
   const removeAccount = React.useCallback<SessionApiContext['removeAccount']>(
     account => {
+      cancelPendingTask()
       dispatch({
         type: 'removed-account',
         accountDid: account.did,
       })
     },
-    [],
+    [cancelPendingTask],
   )
 
   const updateCurrentAccount = React.useCallback<
@@ -422,14 +183,14 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       )
       if (syncedAccount && syncedAccount.refreshJwt) {
         if (syncedAccount.did !== state.currentAgentState.did) {
-          initSession(syncedAccount)
+          resumeSession(syncedAccount)
         } else {
           // @ts-ignore we checked for `refreshJwt` above
           state.currentAgentState.agent.session = syncedAccount
         }
       }
     })
-  }, [state, initSession])
+  }, [state, resumeSession])
 
   const stateContext = React.useMemo(
     () => ({
@@ -447,7 +208,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       createAccount,
       login,
       logout,
-      initSession,
+      resumeSession,
       removeAccount,
       updateCurrentAccount,
     }),
@@ -455,7 +216,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       createAccount,
       login,
       logout,
-      initSession,
+      resumeSession,
       removeAccount,
       updateCurrentAccount,
     ],
@@ -464,8 +225,9 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   // @ts-ignore
   if (IS_DEV && isWeb) window.agent = state.currentAgentState.agent
 
+  const agent = state.currentAgentState.agent as BskyAgent
   return (
-    <AgentContext.Provider value={state.currentAgentState.agent}>
+    <AgentContext.Provider value={agent}>
       <StateContext.Provider value={stateContext}>
         <ApiContext.Provider value={api}>{children}</ApiContext.Provider>
       </StateContext.Provider>
@@ -473,6 +235,18 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   )
 }
 
+function useOneTaskAtATime() {
+  const abortController = React.useRef<AbortController | null>(null)
+  const cancelPendingTask = React.useCallback(() => {
+    if (abortController.current) {
+      abortController.current.abort()
+    }
+    abortController.current = new AbortController()
+    return abortController.current.signal
+  }, [])
+  return cancelPendingTask
+}
+
 export function useSession() {
   return React.useContext(StateContext)
 }

src/state/session/moderation.ts

@@ -0,0 +1,50 @@
+import {BSKY_LABELER_DID, BskyAgent} from '@atproto/api'
+
+import {IS_TEST_USER} from '#/lib/constants'
+import {readLabelers} from './agent-config'
+import {SessionAccount} from './types'
+
+export function configureModerationForGuest() {
+  // This global mutation is *only* OK because this code is only relevant for testing.
+  // Don't add any other global behavior here!
+  switchToBskyAppLabeler()
+}
+
+export async function configureModerationForAccount(
+  agent: BskyAgent,
+  account: SessionAccount,
+) {
+  // This global mutation is *only* OK because this code is only relevant for testing.
+  // Don't add any other global behavior here!
+  switchToBskyAppLabeler()
+  if (IS_TEST_USER(account.handle)) {
+    await trySwitchToTestAppLabeler(agent)
+  }
+
+  // The code below is actually relevant to production (and isn't global).
+  const labelerDids = await readLabelers(account.did).catch(_ => {})
+  if (labelerDids) {
+    agent.configureLabelersHeader(
+      labelerDids.filter(did => did !== BSKY_LABELER_DID),
+    )
+  } else {
+    // If there are no headers in the storage, we'll not send them on the initial requests.
+    // If we wanted to fix this, we could block on the preferences query here.
+  }
+}
+
+function switchToBskyAppLabeler() {
+  BskyAgent.configure({appLabelers: [BSKY_LABELER_DID]})
+}
+
+async function trySwitchToTestAppLabeler(agent: BskyAgent) {
+  const did = (
+    await agent
+      .resolveHandle({handle: 'mod-authority.test'})
+      .catch(_ => undefined)
+  )?.data.did
+  if (did) {
+    console.warn('USING TEST ENV MODERATION')
+    BskyAgent.configure({appLabelers: [did]})
+  }
+}

src/state/session/reducer.ts

@@ -0,0 +1,188 @@
+import {AtpSessionEvent} from '@atproto/api'
+
+import {createPublicAgent} from './agent'
+import {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.
+type OpaqueBskyAgent = {
+  readonly api: unknown
+  readonly app: unknown
+  readonly com: unknown
+}
+
+type AgentState = {
+  readonly agent: OpaqueBskyAgent
+  readonly did: string | undefined
+}
+
+export type State = {
+  readonly accounts: SessionAccount[]
+  readonly currentAgentState: AgentState
+  needsPersist: boolean // Mutated in an effect.
+}
+
+export type Action =
+  | {
+      type: 'received-agent-event'
+      agent: OpaqueBskyAgent
+      accountDid: string
+      refreshedAccount: SessionAccount | undefined
+      sessionEvent: AtpSessionEvent
+    }
+  | {
+      type: 'switched-to-account'
+      newAgent: OpaqueBskyAgent
+      newAccount: SessionAccount
+    }
+  | {
+      type: 'updated-current-account'
+      updatedFields: Partial<
+        Pick<
+          SessionAccount,
+          'handle' | 'email' | 'emailConfirmed' | 'emailAuthFactor'
+        >
+      >
+    }
+  | {
+      type: 'removed-account'
+      accountDid: string
+    }
+  | {
+      type: 'logged-out'
+    }
+  | {
+      type: 'synced-accounts'
+      syncedAccounts: SessionAccount[]
+      syncedCurrentDid: string | undefined
+    }
+
+function createPublicAgentState(): AgentState {
+  return {
+    agent: createPublicAgent(),
+    did: undefined,
+  }
+}
+
+export function getInitialState(persistedAccounts: SessionAccount[]): State {
+  return {
+    accounts: persistedAccounts,
+    currentAgentState: createPublicAgentState(),
+    needsPersist: false,
+  }
+}
+
+export function reducer(state: State, action: Action): State {
+  switch (action.type) {
+    case 'received-agent-event': {
+      const {agent, accountDid, refreshedAccount, sessionEvent} = action
+      if (agent !== state.currentAgentState.agent) {
+        // Only consider events from the active agent.
+        return state
+      }
+      if (sessionEvent === 'network-error') {
+        // Don't change stored accounts but kick to the choose account screen.
+        return {
+          accounts: state.accounts,
+          currentAgentState: createPublicAgentState(),
+          needsPersist: true,
+        }
+      }
+      const existingAccount = state.accounts.find(a => a.did === accountDid)
+      if (
+        !existingAccount ||
+        JSON.stringify(existingAccount) === JSON.stringify(refreshedAccount)
+      ) {
+        // Fast path without a state update.
+        return state
+      }
+      return {
+        accounts: state.accounts.map(a => {
+          if (a.did === accountDid) {
+            if (refreshedAccount) {
+              return refreshedAccount
+            } else {
+              return {
+                ...a,
+                // If we didn't receive a refreshed account, clear out the tokens.
+                accessJwt: undefined,
+                refreshJwt: undefined,
+              }
+            }
+          } else {
+            return a
+          }
+        }),
+        currentAgentState: refreshedAccount
+          ? state.currentAgentState
+          : createPublicAgentState(), // Log out if expired.
+        needsPersist: true,
+      }
+    }
+    case 'switched-to-account': {
+      const {newAccount, newAgent} = action
+      return {
+        accounts: [
+          newAccount,
+          ...state.accounts.filter(a => a.did !== newAccount.did),
+        ],
+        currentAgentState: {
+          did: newAccount.did,
+          agent: newAgent,
+        },
+        needsPersist: true,
+      }
+    }
+    case 'updated-current-account': {
+      const {updatedFields} = action
+      return {
+        accounts: state.accounts.map(a => {
+          if (a.did === state.currentAgentState.did) {
+            return {
+              ...a,
+              ...updatedFields,
+            }
+          } else {
+            return a
+          }
+        }),
+        currentAgentState: state.currentAgentState,
+        needsPersist: true,
+      }
+    }
+    case 'removed-account': {
+      const {accountDid} = action
+      return {
+        accounts: state.accounts.filter(a => a.did !== accountDid),
+        currentAgentState:
+          state.currentAgentState.did === accountDid
+            ? createPublicAgentState() // Log out if removing the current one.
+            : state.currentAgentState,
+        needsPersist: true,
+      }
+    }
+    case 'logged-out': {
+      return {
+        accounts: state.accounts.map(a => ({
+          ...a,
+          // Clear tokens for *every* account (this is a hard logout).
+          refreshJwt: undefined,
+          accessJwt: undefined,
+        })),
+        currentAgentState: createPublicAgentState(),
+        needsPersist: true,
+      }
+    }
+    case 'synced-accounts': {
+      const {syncedAccounts, syncedCurrentDid} = action
+      return {
+        accounts: syncedAccounts,
+        currentAgentState:
+          syncedCurrentDid === state.currentAgentState.did
+            ? state.currentAgentState
+            : createPublicAgentState(), // Log out if different user.
+        needsPersist: false, // Synced from another tab. Don't persist to avoid cycles.
+      }
+    }
+  }
+}

src/state/session/types.ts

@@ -8,6 +8,7 @@ export type SessionStateContext = {
   currentAccount: SessionAccount | undefined
   hasSession: boolean
 }
+
 export type SessionApiContext = {
   createAccount: (props: {
     service: string
@@ -33,10 +34,8 @@ export type SessionApiContext = {
    * access tokens from all accounts, so that returning as any user will
    * require a full login.
    */
-  logout: (
+  logout: (logContext: LogEvents['account:loggedOut']['logContext']) => void
-    logContext: LogEvents['account:loggedOut']['logContext'],
+  resumeSession: (account: SessionAccount) => Promise<void>
-  ) => Promise<void>
-  initSession: (account: SessionAccount) => Promise<void>
   removeAccount: (account: SessionAccount) => void
   updateCurrentAccount: (
     account: Partial<

src/state/session/util.ts

@@ -0,0 +1,36 @@
+import {jwtDecode} from 'jwt-decode'
+
+import {hasProp} from '#/lib/type-guards'
+import {logger} from '#/logger'
+import * as persisted from '#/state/persisted'
+import {SessionAccount} from './types'
+
+export function readLastActiveAccount() {
+  const {currentAccount, accounts} = persisted.get('session')
+  return accounts.find(a => a.did === currentAccount?.did)
+}
+
+export function isSessionDeactivated(accessJwt: string | undefined) {
+  if (accessJwt) {
+    const sessData = jwtDecode(accessJwt)
+    return (
+      hasProp(sessData, 'scope') && sessData.scope === 'com.atproto.deactivated'
+    )
+  }
+  return false
+}
+
+export function isSessionExpired(account: SessionAccount) {
+  try {
+    if (account.accessJwt) {
+      const decoded = jwtDecode(account.accessJwt)
+      if (decoded.exp) {
+        const didExpire = Date.now() >= decoded.exp * 1000
+        return didExpire
+      }
+    }
+  } catch (e) {
+    logger.error(`session: could not decode jwt`)
+  }
+  return true
+}

src/state/session/util/index.ts

@@ -1,186 +0,0 @@
-import {BSKY_LABELER_DID, BskyAgent} from '@atproto/api'
-import {jwtDecode} from 'jwt-decode'
-
-import {IS_PROD_SERVICE, IS_TEST_USER} from '#/lib/constants'
-import {tryFetchGates} from '#/lib/statsig/statsig'
-import {hasProp} from '#/lib/type-guards'
-import {logger} from '#/logger'
-import * as persisted from '#/state/persisted'
-import {DEFAULT_PROD_FEEDS} from '#/state/queries/preferences'
-import {readLabelers} from '../agent-config'
-import {SessionAccount, SessionApiContext} from '../types'
-
-export function isSessionDeactivated(accessJwt: string | undefined) {
-  if (accessJwt) {
-    const sessData = jwtDecode(accessJwt)
-    return (
-      hasProp(sessData, 'scope') && sessData.scope === 'com.atproto.deactivated'
-    )
-  }
-  return false
-}
-
-export function readLastActiveAccount() {
-  const {currentAccount, accounts} = persisted.get('session')
-  return accounts.find(a => a.did === currentAccount?.did)
-}
-
-export function agentToSessionAccount(
-  agent: BskyAgent,
-): SessionAccount | undefined {
-  if (!agent.session) return undefined
-
-  return {
-    service: agent.service.toString(),
-    did: agent.session.did,
-    handle: agent.session.handle,
-    email: agent.session.email,
-    emailConfirmed: agent.session.emailConfirmed || false,
-    emailAuthFactor: agent.session.emailAuthFactor || false,
-    refreshJwt: agent.session.refreshJwt,
-    accessJwt: agent.session.accessJwt,
-    deactivated: isSessionDeactivated(agent.session.accessJwt),
-    pdsUrl: agent.pdsUrl?.toString(),
-  }
-}
-
-export function configureModerationForGuest() {
-  switchToBskyAppLabeler()
-}
-
-export async function configureModerationForAccount(
-  agent: BskyAgent,
-  account: SessionAccount,
-) {
-  switchToBskyAppLabeler()
-  if (IS_TEST_USER(account.handle)) {
-    await trySwitchToTestAppLabeler(agent)
-  }
-
-  const labelerDids = await readLabelers(account.did).catch(_ => {})
-  if (labelerDids) {
-    agent.configureLabelersHeader(
-      labelerDids.filter(did => did !== BSKY_LABELER_DID),
-    )
-  } else {
-    // If there are no headers in the storage, we'll not send them on the initial requests.
-    // If we wanted to fix this, we could block on the preferences query here.
-  }
-}
-
-function switchToBskyAppLabeler() {
-  BskyAgent.configure({appLabelers: [BSKY_LABELER_DID]})
-}
-
-async function trySwitchToTestAppLabeler(agent: BskyAgent) {
-  const did = (
-    await agent
-      .resolveHandle({handle: 'mod-authority.test'})
-      .catch(_ => undefined)
-  )?.data.did
-  if (did) {
-    console.warn('USING TEST ENV MODERATION')
-    BskyAgent.configure({appLabelers: [did]})
-  }
-}
-
-export function isSessionExpired(account: SessionAccount) {
-  try {
-    if (account.accessJwt) {
-      const decoded = jwtDecode(account.accessJwt)
-      if (decoded.exp) {
-        const didExpire = Date.now() >= decoded.exp * 1000
-        return didExpire
-      }
-    }
-  } catch (e) {
-    logger.error(`session: could not decode jwt`)
-  }
-  return true
-}
-
-export async function createAgentAndLogin({
-  service,
-  identifier,
-  password,
-  authFactorToken,
-}: {
-  service: string
-  identifier: string
-  password: string
-  authFactorToken?: string
-}) {
-  const agent = new BskyAgent({service})
-  await agent.login({identifier, password, authFactorToken})
-
-  const account = agentToSessionAccount(agent)
-  if (!agent.session || !account) {
-    throw new Error(`session: login failed to establish a session`)
-  }
-
-  const fetchingGates = tryFetchGates(account.did, 'prefer-fresh-gates')
-  await configureModerationForAccount(agent, account)
-
-  return {
-    agent,
-    account,
-    fetchingGates,
-  }
-}
-
-export async function createAgentAndCreateAccount({
-  service,
-  email,
-  password,
-  handle,
-  birthDate,
-  inviteCode,
-  verificationPhone,
-  verificationCode,
-}: Parameters<SessionApiContext['createAccount']>[0]) {
-  const agent = new BskyAgent({service})
-  await agent.createAccount({
-    email,
-    password,
-    handle,
-    inviteCode,
-    verificationPhone,
-    verificationCode,
-  })
-
-  const account = agentToSessionAccount(agent)!
-  if (!agent.session || !account) {
-    throw new Error(`session: createAccount failed to establish a session`)
-  }
-
-  const fetchingGates = tryFetchGates(account.did, 'prefer-fresh-gates')
-
-  if (!account.deactivated) {
-    /*dont await*/ agent.upsertProfile(_existing => {
-      return {
-        displayName: '',
-
-        // HACKFIX
-        // creating a bunch of identical profile objects is breaking the relay
-        // tossing this unspecced field onto it to reduce the size of the problem
-        // -prf
-        createdAt: new Date().toISOString(),
-      }
-    })
-  }
-
-  // Not awaited so that we can still get into onboarding.
-  // This is OK because we won't let you toggle adult stuff until you set the date.
-  agent.setPersonalDetails({birthDate: birthDate.toISOString()})
-  if (IS_PROD_SERVICE(service)) {
-    agent.setSavedFeeds(DEFAULT_PROD_FEEDS.saved, DEFAULT_PROD_FEEDS.pinned)
-  }
-
-  await configureModerationForAccount(agent, account)
-
-  return {
-    agent,
-    account,
-    fetchingGates,
-  }
-}