Session fixes, pt. 1 (#3762)

* Update persisted schema for new source of truth, implement in existing session

(cherry picked from commit b1e5f12baee932721d66c60dd51c981b46b0c274)

* Improve toasts, log caught error, during switch account

(cherry picked from commit fe0d1507063d2e532b7b1a447670b689292d1dc3)

* Handle thrown errors from initSession during login

(cherry picked from commit 2c85c045917e923901284b9ba310a82e28f37b5c)

---------

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

4 files changed, 62 insertions, 24 deletions

diff --git a/src/lib/hooks/useAccountSwitcher.ts b/src/lib/hooks/useAccountSwitcher.ts
index 6a1cea234..6d2f7b36b 100644
--- a/src/lib/hooks/useAccountSwitcher.ts
+++ b/src/lib/hooks/useAccountSwitcher.ts
@@ -1,6 +1,9 @@
 import {useCallback} from 'react'
+import {msg} from '@lingui/macro'
+import {useLingui} from '@lingui/react'
 
 import {useAnalytics} from '#/lib/analytics/analytics'
+import {logger} from '#/logger'
 import {isWeb} from '#/platform/detection'
 import {SessionAccount, useSessionApi} from '#/state/session'
 import {useLoggedOutViewControls} from '#/state/shell/logged-out'
@@ -8,6 +11,7 @@ import * as Toast from '#/view/com/util/Toast'
 import {LogEvents} from '../statsig/statsig'
 
 export function useAccountSwitcher() {
+  const {_} = useLingui()
   const {track} = useAnalytics()
   const {selectAccount, clearCurrentAccount} = useSessionApi()
   const {requestSwitchToAccount} = useLoggedOutViewControls()
@@ -31,21 +35,26 @@ export function useAccountSwitcher() {
           }
           await selectAccount(account, logContext)
           setTimeout(() => {
-            Toast.show(`Signed in as @${account.handle}`)
+            Toast.show(_(msg`Signed in as @${account.handle}`))
           }, 100)
         } else {
           requestSwitchToAccount({requestedAccount: account.did})
           Toast.show(
-            `Please sign in as @${account.handle}`,
+            _(msg`Please sign in as @${account.handle}`),
             'circle-exclamation',
           )
         }
-      } catch (e) {
-        Toast.show('Sorry! We need you to enter your password.')
+      } catch (e: any) {
+        logger.error(`switch account: selectAccount failed`, {
+          message: e.message,
+        })
         clearCurrentAccount() // back user out to login
+        setTimeout(() => {
+          Toast.show(_(msg`Sorry! We need you to enter your password.`))
+        }, 100)
       }
     },
-    [track, clearCurrentAccount, selectAccount, requestSwitchToAccount],
+    [_, track, clearCurrentAccount, selectAccount, requestSwitchToAccount],
   )
 
   return {onPressSwitchAccount}
diff --git a/src/screens/Login/ChooseAccountForm.tsx b/src/screens/Login/ChooseAccountForm.tsx
index 134411903..d36d10977 100644
--- a/src/screens/Login/ChooseAccountForm.tsx
+++ b/src/screens/Login/ChooseAccountForm.tsx
@@ -5,6 +5,7 @@ import {useLingui} from '@lingui/react'
 
 import {useAnalytics} from '#/lib/analytics/analytics'
 import {logEvent} from '#/lib/statsig/statsig'
+import {logger} from '#/logger'
 import {SessionAccount, useSession, useSessionApi} from '#/state/session'
 import {useLoggedOutViewControls} from '#/state/shell/logged-out'
 import * as Toast from '#/view/com/util/Toast'
@@ -38,15 +39,22 @@ export const ChooseAccountForm = ({
           setShowLoggedOut(false)
           Toast.show(_(msg`Already signed in as @${account.handle}`))
         } else {
-          await initSession(account)
-          logEvent('account:loggedIn', {
-            logContext: 'ChooseAccountForm',
-            withPassword: false,
-          })
-          track('Sign In', {resumedSession: true})
-          setTimeout(() => {
-            Toast.show(_(msg`Signed in as @${account.handle}`))
-          }, 100)
+          try {
+            await initSession(account)
+            logEvent('account:loggedIn', {
+              logContext: 'ChooseAccountForm',
+              withPassword: false,
+            })
+            track('Sign In', {resumedSession: true})
+            setTimeout(() => {
+              Toast.show(_(msg`Signed in as @${account.handle}`))
+            }, 100)
+          } catch (e: any) {
+            logger.error('choose account: initSession failed', {
+              message: e.message,
+            })
+            onSelectAccount(account)
+          }
         }
       } else {
         onSelectAccount(account)
diff --git a/src/state/persisted/schema.ts b/src/state/persisted/schema.ts
index f090365a3..43f08e12e 100644
--- a/src/state/persisted/schema.ts
+++ b/src/state/persisted/schema.ts
@@ -4,7 +4,10 @@ import {deviceLocales, prefersReducedMotion} from '#/platform/detection'
 
 const externalEmbedOptions = ['show', 'hide'] as const
 
-// only data needed for rendering account page
+/**
+ * A account persisted to storage. Stored in the `accounts[]` array. Contains
+ * base account info and access tokens.
+ */
 const accountSchema = z.object({
   service: z.string(),
   did: z.string(),
@@ -19,12 +22,26 @@ const accountSchema = z.object({
 })
 export type PersistedAccount = z.infer<typeof accountSchema>
 
+/**
+ * The current account. Stored in the `currentAccount` field.
+ *
+ * In previous versions, this included tokens and other info. Now, it's used
+ * only to reference the `did` field, and all other fields are marked as
+ * optional. They should be considered deprecated and not used, but are kept
+ * here for backwards compat.
+ */
+const currentAccountSchema = accountSchema.extend({
+  service: z.string().optional(),
+  handle: z.string().optional(),
+})
+export type PersistedCurrentAccount = z.infer<typeof currentAccountSchema>
+
 export const schema = z.object({
   colorMode: z.enum(['system', 'light', 'dark']),
   darkTheme: z.enum(['dim', 'dark']).optional(),
   session: z.object({
     accounts: z.array(accountSchema),
-    currentAccount: accountSchema.optional(),
+    currentAccount: currentAccountSchema.optional(),
   }),
   reminders: z.object({
     lastEmailConfirm: z.string().optional(),
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index e45aa031f..37b108766 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -618,20 +618,24 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
       logger.debug(`session: persisted onUpdate`, {})
 
-      if (session.currentAccount && session.currentAccount.refreshJwt) {
-        if (session.currentAccount?.did !== state.currentAccount?.did) {
+      const selectedAccount = session.accounts.find(
+        a => a.did === session.currentAccount?.did,
+      )
+
+      if (selectedAccount && selectedAccount.refreshJwt) {
+        if (selectedAccount.did !== state.currentAccount?.did) {
           logger.debug(`session: persisted onUpdate, switching accounts`, {
             from: {
               did: state.currentAccount?.did,
               handle: state.currentAccount?.handle,
             },
             to: {
-              did: session.currentAccount.did,
-              handle: session.currentAccount.handle,
+              did: selectedAccount.did,
+              handle: selectedAccount.handle,
             },
           })
 
-          initSession(session.currentAccount)
+          initSession(selectedAccount)
         } else {
           logger.debug(`session: persisted onUpdate, updating session`, {})
 
@@ -641,9 +645,9 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
            * already persisted, and we'll get a loop between tabs.
            */
           // @ts-ignore we checked for `refreshJwt` above
-          __globalAgent.session = session.currentAccount
+          __globalAgent.session = selectedAccount
         }
-      } else if (!session.currentAccount && state.currentAccount) {
+      } else if (!selectedAccount && state.currentAccount) {
         logger.debug(
           `session: persisted onUpdate, logging out`,
           {},
@@ -662,7 +666,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       setState(s => ({
         ...s,
         accounts: session.accounts,
-        currentAccount: session.currentAccount,
+        currentAccount: selectedAccount,
       }))
     })
   }, [state, setState, clearCurrentAccount, initSession])