First pass at a session handler (#1850)

* First pass at a session handler

* TODOs

* Fix recursion

* Couple more things

* Add back resume session concept

* Handle ready

* Cleanup of initial loading states

* Handle init failure

* Cleanup

* Remove account

* Add updateCurrentAccount

* Remove log

* Cleanup

* Integrate removeAccount

* Add hasSession

* Add to App.native, harden migration

* Use effect to persist data

9 files changed, 488 insertions, 56 deletions

diff --git a/src/App.native.tsx b/src/App.native.tsx
index 9f38d133d..8479465fd 100644
--- a/src/App.native.tsx
+++ b/src/App.native.tsx
@@ -26,6 +26,12 @@ import {Provider as ModalStateProvider} from 'state/modals'
 import {Provider as MutedThreadsProvider} from 'state/muted-threads'
 import {Provider as InvitesStateProvider} from 'state/invites'
 import {Provider as PrefsStateProvider} from 'state/preferences'
+import {
+  Provider as SessionProvider,
+  useSession,
+  useSessionApi,
+} from 'state/session'
+import * as persisted from '#/state/persisted'
 import {i18n} from '@lingui/core'
 import {I18nProvider} from '@lingui/react'
 import {messages} from './locale/locales/en/messages'
@@ -36,6 +42,8 @@ SplashScreen.preventAutoHideAsync()
 
 const InnerApp = observer(function AppImpl() {
   const colorMode = useColorMode()
+  const {isInitialLoad} = useSession()
+  const {resumeSession} = useSessionApi()
   const [rootStore, setRootStore] = useState<RootStoreModel | undefined>(
     undefined,
   )
@@ -52,10 +60,17 @@ const InnerApp = observer(function AppImpl() {
     })
   }, [])
 
+  useEffect(() => {
+    const account = persisted.get('session').currentAccount
+    resumeSession(account)
+  }, [resumeSession])
+
   // show nothing prior to init
-  if (!rootStore) {
+  if (!rootStore || isInitialLoad) {
+    // TODO add a loading state
     return null
   }
+
   return (
     <QueryClientProvider client={queryClient}>
       <ThemeProvider theme={colorMode}>
@@ -88,17 +103,19 @@ function App() {
   }
 
   return (
-    <ShellStateProvider>
-      <PrefsStateProvider>
-        <MutedThreadsProvider>
-          <InvitesStateProvider>
-            <ModalStateProvider>
-              <InnerApp />
-            </ModalStateProvider>
-          </InvitesStateProvider>
-        </MutedThreadsProvider>
-      </PrefsStateProvider>
-    </ShellStateProvider>
+    <SessionProvider>
+      <ShellStateProvider>
+        <PrefsStateProvider>
+          <MutedThreadsProvider>
+            <InvitesStateProvider>
+              <ModalStateProvider>
+                <InnerApp />
+              </ModalStateProvider>
+            </InvitesStateProvider>
+          </MutedThreadsProvider>
+        </PrefsStateProvider>
+      </ShellStateProvider>
+    </SessionProvider>
   )
 }
 
diff --git a/src/App.web.tsx b/src/App.web.tsx
index ef275b392..fc76afce1 100644
--- a/src/App.web.tsx
+++ b/src/App.web.tsx
@@ -24,8 +24,16 @@ import {Provider as ModalStateProvider} from 'state/modals'
 import {Provider as MutedThreadsProvider} from 'state/muted-threads'
 import {Provider as InvitesStateProvider} from 'state/invites'
 import {Provider as PrefsStateProvider} from 'state/preferences'
+import {
+  Provider as SessionProvider,
+  useSession,
+  useSessionApi,
+} from 'state/session'
+import * as persisted from '#/state/persisted'
 
 const InnerApp = observer(function AppImpl() {
+  const {isInitialLoad} = useSession()
+  const {resumeSession} = useSessionApi()
   const colorMode = useColorMode()
   const [rootStore, setRootStore] = useState<RootStoreModel | undefined>(
     undefined,
@@ -38,10 +46,16 @@ const InnerApp = observer(function AppImpl() {
       analytics.init(store)
     })
     dynamicActivate(defaultLocale) // async import of locale data
-  }, [])
+  }, [resumeSession])
+
+  useEffect(() => {
+    const account = persisted.get('session').currentAccount
+    resumeSession(account)
+  }, [resumeSession])
 
   // show nothing prior to init
-  if (!rootStore) {
+  if (!rootStore || isInitialLoad) {
+    // TODO add a loading state
     return null
   }
 
@@ -77,17 +91,19 @@ function App() {
   }
 
   return (
-    <ShellStateProvider>
-      <PrefsStateProvider>
-        <MutedThreadsProvider>
-          <InvitesStateProvider>
-            <ModalStateProvider>
-              <InnerApp />
-            </ModalStateProvider>
-          </InvitesStateProvider>
-        </MutedThreadsProvider>
-      </PrefsStateProvider>
-    </ShellStateProvider>
+    <SessionProvider>
+      <ShellStateProvider>
+        <PrefsStateProvider>
+          <MutedThreadsProvider>
+            <InvitesStateProvider>
+              <ModalStateProvider>
+                <InnerApp />
+              </ModalStateProvider>
+            </InvitesStateProvider>
+          </MutedThreadsProvider>
+        </PrefsStateProvider>
+      </ShellStateProvider>
+    </SessionProvider>
   )
 }
 
diff --git a/src/state/persisted/index.ts b/src/state/persisted/index.ts
index f43cc9527..a5c38513f 100644
--- a/src/state/persisted/index.ts
+++ b/src/state/persisted/index.ts
@@ -5,7 +5,7 @@ import {migrate} from '#/state/persisted/legacy'
 import * as store from '#/state/persisted/store'
 import BroadcastChannel from '#/state/persisted/broadcast'
 
-export type {Schema} from '#/state/persisted/schema'
+export type {Schema, PersistedAccount} from '#/state/persisted/schema'
 export {defaults} from '#/state/persisted/schema'
 
 const broadcast = new BroadcastChannel('BSKY_BROADCAST_CHANNEL')
@@ -50,7 +50,9 @@ export async function write<K extends keyof Schema>(
     await store.write(_state)
     // must happen on next tick, otherwise the tab will read stale storage data
     setTimeout(() => broadcast.postMessage({event: UPDATE_EVENT}), 0)
-    logger.debug(`persisted state: wrote root state to storage`)
+    logger.debug(`persisted state: wrote root state to storage`, {
+      updatedKey: key,
+    })
   } catch (e) {
     logger.error(`persisted state: failed writing root state to storage`, {
       error: e,
diff --git a/src/state/persisted/legacy.ts b/src/state/persisted/legacy.ts
index 3da509304..fefa7f372 100644
--- a/src/state/persisted/legacy.ts
+++ b/src/state/persisted/legacy.ts
@@ -66,43 +66,45 @@ type LegacySchema = {
 
 const DEPRECATED_ROOT_STATE_STORAGE_KEY = 'root'
 
-export function transform(legacy: LegacySchema): Schema {
+// TODO remove, assume that partial data may be here during our refactor
+export function transform(legacy: Partial<LegacySchema>): Schema {
   return {
     colorMode: legacy.shell?.colorMode || defaults.colorMode,
     session: {
-      accounts: legacy.session.accounts || defaults.session.accounts,
+      accounts: legacy.session?.accounts || defaults.session.accounts,
       currentAccount:
-        legacy.session.accounts.find(a => a.did === legacy.session.data.did) ||
-        defaults.session.currentAccount,
+        legacy.session?.accounts?.find(
+          a => a.did === legacy.session?.data?.did,
+        ) || defaults.session.currentAccount,
     },
     reminders: {
       lastEmailConfirm:
-        legacy.reminders.lastEmailConfirm ||
+        legacy.reminders?.lastEmailConfirm ||
         defaults.reminders.lastEmailConfirm,
     },
     languagePrefs: {
       primaryLanguage:
-        legacy.preferences.primaryLanguage ||
+        legacy.preferences?.primaryLanguage ||
         defaults.languagePrefs.primaryLanguage,
       contentLanguages:
-        legacy.preferences.contentLanguages ||
+        legacy.preferences?.contentLanguages ||
         defaults.languagePrefs.contentLanguages,
       postLanguage:
-        legacy.preferences.postLanguage || defaults.languagePrefs.postLanguage,
+        legacy.preferences?.postLanguage || defaults.languagePrefs.postLanguage,
       postLanguageHistory:
-        legacy.preferences.postLanguageHistory ||
+        legacy.preferences?.postLanguageHistory ||
         defaults.languagePrefs.postLanguageHistory,
     },
     requireAltTextEnabled:
-      legacy.preferences.requireAltTextEnabled ||
+      legacy.preferences?.requireAltTextEnabled ||
       defaults.requireAltTextEnabled,
-    mutedThreads: legacy.mutedThreads.uris || defaults.mutedThreads,
+    mutedThreads: legacy.mutedThreads?.uris || defaults.mutedThreads,
     invites: {
       copiedInvites:
-        legacy.invitedUsers.copiedInvites || defaults.invites.copiedInvites,
+        legacy.invitedUsers?.copiedInvites || defaults.invites.copiedInvites,
     },
     onboarding: {
-      step: legacy.onboarding.step || defaults.onboarding.step,
+      step: legacy.onboarding?.step || defaults.onboarding.step,
     },
   }
 }
diff --git a/src/state/persisted/schema.ts b/src/state/persisted/schema.ts
index 9c52661e4..a510262fb 100644
--- a/src/state/persisted/schema.ts
+++ b/src/state/persisted/schema.ts
@@ -2,15 +2,17 @@ import {z} from 'zod'
 import {deviceLocales} from '#/platform/detection'
 
 // only data needed for rendering account page
+// TODO agent.resumeSession requires the following fields
 const accountSchema = z.object({
   service: z.string(),
   did: z.string(),
-  refreshJwt: z.string().optional(),
-  accessJwt: z.string().optional(),
-  handle: z.string().optional(),
-  displayName: z.string().optional(),
-  aviUrl: z.string().optional(),
+  handle: z.string(),
+  refreshJwt: z.string().optional(), // optional because it can expire
+  accessJwt: z.string().optional(), // optional because it can expire
+  // displayName: z.string().optional(),
+  // aviUrl: z.string().optional(),
 })
+export type PersistedAccount = z.infer<typeof accountSchema>
 
 export const schema = z.object({
   colorMode: z.enum(['system', 'light', 'dark']),
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
new file mode 100644
index 000000000..0f3118168
--- /dev/null
+++ b/src/state/session/index.tsx
@@ -0,0 +1,384 @@
+import React from 'react'
+import {BskyAgent, AtpPersistSessionHandler} from '@atproto/api'
+
+import {networkRetry} from '#/lib/async/retry'
+import {logger} from '#/logger'
+import * as persisted from '#/state/persisted'
+
+export type SessionAccount = persisted.PersistedAccount
+
+export type StateContext = {
+  isInitialLoad: boolean
+  agent: BskyAgent
+  accounts: persisted.PersistedAccount[]
+  currentAccount: persisted.PersistedAccount | undefined
+  hasSession: boolean
+}
+export type ApiContext = {
+  createAccount: (props: {
+    service: string
+    email: string
+    password: string
+    handle: string
+    inviteCode?: string
+  }) => Promise<void>
+  login: (props: {
+    service: string
+    identifier: string
+    password: string
+  }) => Promise<void>
+  logout: () => Promise<void>
+  initSession: (account: persisted.PersistedAccount) => Promise<void>
+  resumeSession: (account?: persisted.PersistedAccount) => Promise<void>
+  removeAccount: (
+    account: Partial<Pick<persisted.PersistedAccount, 'handle' | 'did'>>,
+  ) => void
+  updateCurrentAccount: (
+    account: Pick<persisted.PersistedAccount, 'handle'>,
+  ) => void
+}
+
+export const PUBLIC_BSKY_AGENT = new BskyAgent({
+  service: 'https://api.bsky.app',
+})
+
+const StateContext = React.createContext<StateContext>({
+  hasSession: false,
+  isInitialLoad: true,
+  accounts: [],
+  currentAccount: undefined,
+  agent: PUBLIC_BSKY_AGENT,
+})
+
+const ApiContext = React.createContext<ApiContext>({
+  createAccount: async () => {},
+  login: async () => {},
+  logout: async () => {},
+  initSession: async () => {},
+  resumeSession: async () => {},
+  removeAccount: () => {},
+  updateCurrentAccount: () => {},
+})
+
+function createPersistSessionHandler(
+  account: persisted.PersistedAccount,
+  persistSessionCallback: (props: {
+    expired: boolean
+    refreshedAccount: persisted.PersistedAccount
+  }) => void,
+): AtpPersistSessionHandler {
+  return function persistSession(event, session) {
+    const expired = !(event === 'create' || event === 'update')
+    const refreshedAccount = {
+      service: account.service,
+      did: session?.did || account.did,
+      handle: session?.handle || account.handle,
+      refreshJwt: session?.refreshJwt, // undefined when expired or creation fails
+      accessJwt: session?.accessJwt, // undefined when expired or creation fails
+    }
+
+    logger.debug(`session: BskyAgent.persistSession`, {
+      expired,
+      did: refreshedAccount.did,
+      handle: refreshedAccount.handle,
+    })
+
+    persistSessionCallback({
+      expired,
+      refreshedAccount,
+    })
+  }
+}
+
+export function Provider({children}: React.PropsWithChildren<{}>) {
+  const [state, setState] = React.useState<StateContext>({
+    hasSession: false,
+    isInitialLoad: true, // try to resume the session first
+    accounts: persisted.get('session').accounts,
+    currentAccount: undefined, // assume logged out to start
+    agent: PUBLIC_BSKY_AGENT,
+  })
+
+  const upsertAccount = React.useCallback(
+    (account: persisted.PersistedAccount, expired = false) => {
+      setState(s => {
+        return {
+          ...s,
+          currentAccount: expired ? undefined : account,
+          accounts: [account, ...s.accounts.filter(a => a.did !== account.did)],
+        }
+      })
+    },
+    [setState],
+  )
+
+  // TODO have not connected this yet
+  const createAccount = React.useCallback<ApiContext['createAccount']>(
+    async ({service, email, password, handle, inviteCode}: any) => {
+      logger.debug(`session: creating account`, {
+        service,
+        handle,
+      })
+
+      const agent = new BskyAgent({service})
+
+      await agent.createAccount({
+        handle,
+        password,
+        email,
+        inviteCode,
+      })
+
+      if (!agent.session) {
+        throw new Error(`session: createAccount failed to establish a session`)
+      }
+
+      const account: persisted.PersistedAccount = {
+        service,
+        did: agent.session.did,
+        refreshJwt: agent.session.refreshJwt,
+        accessJwt: agent.session.accessJwt,
+        handle: agent.session.handle,
+      }
+
+      agent.setPersistSessionHandler(
+        createPersistSessionHandler(account, ({expired, refreshedAccount}) => {
+          upsertAccount(refreshedAccount, expired)
+        }),
+      )
+
+      upsertAccount(account)
+
+      logger.debug(`session: created account`, {
+        service,
+        handle,
+      })
+    },
+    [upsertAccount],
+  )
+
+  const login = React.useCallback<ApiContext['login']>(
+    async ({service, identifier, password}) => {
+      logger.debug(`session: login`, {
+        service,
+        identifier,
+      })
+
+      const agent = new BskyAgent({service})
+
+      await agent.login({identifier, password})
+
+      if (!agent.session) {
+        throw new Error(`session: login failed to establish a session`)
+      }
+
+      const account: persisted.PersistedAccount = {
+        service,
+        did: agent.session.did,
+        refreshJwt: agent.session.refreshJwt,
+        accessJwt: agent.session.accessJwt,
+        handle: agent.session.handle,
+      }
+
+      agent.setPersistSessionHandler(
+        createPersistSessionHandler(account, ({expired, refreshedAccount}) => {
+          upsertAccount(refreshedAccount, expired)
+        }),
+      )
+
+      upsertAccount(account)
+
+      logger.debug(`session: logged in`, {
+        service,
+        identifier,
+      })
+    },
+    [upsertAccount],
+  )
+
+  const logout = React.useCallback<ApiContext['logout']>(async () => {
+    logger.debug(`session: logout`)
+    setState(s => {
+      return {
+        ...s,
+        agent: PUBLIC_BSKY_AGENT,
+        currentAccount: undefined,
+        accounts: s.accounts.map(a => ({
+          ...a,
+          refreshJwt: undefined,
+          accessJwt: undefined,
+        })),
+      }
+    })
+  }, [setState])
+
+  const initSession = React.useCallback<ApiContext['initSession']>(
+    async account => {
+      logger.debug(`session: initSession`, {
+        did: account.did,
+        handle: account.handle,
+      })
+
+      const agent = new BskyAgent({
+        service: account.service,
+        persistSession: createPersistSessionHandler(
+          account,
+          ({expired, refreshedAccount}) => {
+            upsertAccount(refreshedAccount, expired)
+          },
+        ),
+      })
+
+      await networkRetry(3, () =>
+        agent.resumeSession({
+          accessJwt: account.accessJwt || '',
+          refreshJwt: account.refreshJwt || '',
+          did: account.did,
+          handle: account.handle,
+        }),
+      )
+
+      upsertAccount(account)
+    },
+    [upsertAccount],
+  )
+
+  const resumeSession = React.useCallback<ApiContext['resumeSession']>(
+    async account => {
+      try {
+        if (account) {
+          await initSession(account)
+        }
+      } catch (e) {
+        logger.error(`session: resumeSession failed`, {error: e})
+      } finally {
+        setState(s => ({
+          ...s,
+          isInitialLoad: false,
+        }))
+      }
+    },
+    [initSession],
+  )
+
+  const removeAccount = React.useCallback<ApiContext['removeAccount']>(
+    account => {
+      setState(s => {
+        return {
+          ...s,
+          accounts: s.accounts.filter(
+            a => !(a.did === account.did || a.handle === account.handle),
+          ),
+        }
+      })
+    },
+    [setState],
+  )
+
+  const updateCurrentAccount = React.useCallback<
+    ApiContext['updateCurrentAccount']
+  >(
+    account => {
+      setState(s => {
+        const currentAccount = s.currentAccount
+
+        // ignore, should never happen
+        if (!currentAccount) return s
+
+        const updatedAccount = {
+          ...currentAccount,
+          handle: account.handle, // only update handle rn
+        }
+
+        return {
+          ...s,
+          currentAccount: updatedAccount,
+          accounts: s.accounts.filter(a => a.did !== currentAccount.did),
+        }
+      })
+    },
+    [setState],
+  )
+
+  React.useEffect(() => {
+    persisted.write('session', {
+      accounts: state.accounts,
+      currentAccount: state.currentAccount,
+    })
+  }, [state])
+
+  React.useEffect(() => {
+    return persisted.onUpdate(() => {
+      const session = persisted.get('session')
+
+      logger.debug(`session: onUpdate`)
+
+      if (session.currentAccount) {
+        if (session.currentAccount?.did !== state.currentAccount?.did) {
+          logger.debug(`session: switching account`, {
+            from: {
+              did: state.currentAccount?.did,
+              handle: state.currentAccount?.handle,
+            },
+            to: {
+              did: session.currentAccount.did,
+              handle: session.currentAccount.handle,
+            },
+          })
+
+          initSession(session.currentAccount)
+        }
+      } else if (!session.currentAccount && state.currentAccount) {
+        logger.debug(`session: logging out`, {
+          did: state.currentAccount?.did,
+          handle: state.currentAccount?.handle,
+        })
+
+        logout()
+      }
+    })
+  }, [state, logout, initSession])
+
+  const stateContext = React.useMemo(
+    () => ({
+      ...state,
+      hasSession: !!state.currentAccount,
+    }),
+    [state],
+  )
+
+  const api = React.useMemo(
+    () => ({
+      createAccount,
+      login,
+      logout,
+      initSession,
+      resumeSession,
+      removeAccount,
+      updateCurrentAccount,
+    }),
+    [
+      createAccount,
+      login,
+      logout,
+      initSession,
+      resumeSession,
+      removeAccount,
+      updateCurrentAccount,
+    ],
+  )
+
+  return (
+    <StateContext.Provider value={stateContext}>
+      <ApiContext.Provider value={api}>{children}</ApiContext.Provider>
+    </StateContext.Provider>
+  )
+}
+
+export function useSession() {
+  return React.useContext(StateContext)
+}
+
+export function useSessionApi() {
+  return React.useContext(ApiContext)
+}
diff --git a/src/state/shell/onboarding.tsx b/src/state/shell/onboarding.tsx
index 5963cc50e..6a18b461f 100644
--- a/src/state/shell/onboarding.tsx
+++ b/src/state/shell/onboarding.tsx
@@ -82,12 +82,16 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
   React.useEffect(() => {
     return persisted.onUpdate(() => {
-      dispatch({
-        type: 'set',
-        step: persisted.get('onboarding').step as OnboardingStep,
-      })
+      const next = persisted.get('onboarding').step
+      // TODO we've introduced a footgun
+      if (state.step !== next) {
+        dispatch({
+          type: 'set',
+          step: persisted.get('onboarding').step as OnboardingStep,
+        })
+      }
     })
-  }, [dispatch])
+  }, [state, dispatch])
 
   return (
     <stateContext.Provider value={state}>
diff --git a/src/view/com/util/AccountDropdownBtn.tsx b/src/view/com/util/AccountDropdownBtn.tsx
index 2042531e9..158ed9b6d 100644
--- a/src/view/com/util/AccountDropdownBtn.tsx
+++ b/src/view/com/util/AccountDropdownBtn.tsx
@@ -5,23 +5,23 @@ import {
   FontAwesomeIconStyle,
 } from '@fortawesome/react-native-fontawesome'
 import {s} from 'lib/styles'
-import {useStores} from 'state/index'
 import {usePalette} from 'lib/hooks/usePalette'
 import {DropdownItem, NativeDropdown} from './forms/NativeDropdown'
 import * as Toast from '../../com/util/Toast'
+import {useSessionApi} from '#/state/session'
 import {useLingui} from '@lingui/react'
 import {msg} from '@lingui/macro'
 
 export function AccountDropdownBtn({handle}: {handle: string}) {
-  const store = useStores()
   const pal = usePalette('default')
+  const {removeAccount} = useSessionApi()
   const {_} = useLingui()
 
   const items: DropdownItem[] = [
     {
       label: 'Remove account',
       onPress: () => {
-        store.session.removeAccount(handle)
+        removeAccount({handle})
         Toast.show('Account removed from quick access')
       },
       icon: {
diff --git a/src/view/screens/Settings.tsx b/src/view/screens/Settings.tsx
index 010de23db..062533c27 100644
--- a/src/view/screens/Settings.tsx
+++ b/src/view/screens/Settings.tsx
@@ -57,6 +57,7 @@ import {
   useRequireAltTextEnabled,
   useSetRequireAltTextEnabled,
 } from '#/state/preferences'
+import {useSession, useSessionApi} from '#/state/session'
 
 // TEMPORARY (APP-700)
 // remove after backend testing finishes
@@ -87,6 +88,8 @@ export const SettingsScreen = withAuthRequired(
       store.agent,
     )
     const {openModal} = useModalControls()
+    const {logout} = useSessionApi()
+    const {accounts} = useSession()
 
     const primaryBg = useCustomPalette<ViewStyle>({
       light: {backgroundColor: colors.blue0},
@@ -153,8 +156,9 @@ export const SettingsScreen = withAuthRequired(
 
     const onPressSignout = React.useCallback(() => {
       track('Settings:SignOutButtonClicked')
+      logout()
       store.session.logout()
-    }, [track, store])
+    }, [track, store, logout])
 
     const onPressDeleteAccount = React.useCallback(() => {
       openModal({name: 'delete-account'})
@@ -294,7 +298,7 @@ export const SettingsScreen = withAuthRequired(
               </View>
             </Link>
           )}
-          {store.session.switchableAccounts.map(account => (
+          {accounts.map(account => (
             <TouchableOpacity
               testID={`switchToAccountBtn-${account.handle}`}
               key={account.did}
@@ -306,10 +310,11 @@ export const SettingsScreen = withAuthRequired(
               accessibilityLabel={`Switch to ${account.handle}`}
               accessibilityHint="Switches the account you are logged in to">
               <View style={styles.avi}>
-                <UserAvatar size={40} avatar={account.aviUrl} />
+                {/*<UserAvatar size={40} avatar={account.aviUrl} />*/}
               </View>
               <View style={[s.flex1]}>
                 <Text type="md-bold" style={pal.text}>
+                  {/* @ts-ignore */}
                   {account.displayName || account.handle}
                 </Text>
                 <Text type="sm" style={pal.textLight}>