Merge pull request #25 from blockworks-foundation/feature/fix-mango-accounts-after-reconnection

mango accounts reaload on wallet change

components/TopBar.tsx

@@ -19,7 +19,7 @@ const TopBar = () => {
   const { t } = useTranslation('common')
   const mangoAccount = mangoStore((s) => s.mangoAccount.current)
   const connected = mangoStore((s) => s.connected)
-  const [isOnboarded] = useLocalStorageState(IS_ONBOARDED_KEY)
+  const [isOnBoarded] = useLocalStorageState(IS_ONBOARDED_KEY)
   const [showUserSetupModal, setShowUserSetupModal] = useState(false)
   const [showCreateAccountModal, setShowCreateAccountModal] = useState(false)
   const [showMangoAccountsModal, setShowMangoAccountsModal] = useState(false)
@@ -91,7 +91,7 @@ const TopBar = () => {
             </button>
             <ConnectedMenu />
           </div>
-        ) : isOnboarded ? (
+        ) : isOnBoarded ? (
           <ConnectWalletButton />
         ) : (
           <button

components/account/AccountPage.tsx

@@ -73,7 +73,7 @@ const AccountPage = () => {
   const [showExpandChart, setShowExpandChart] = useState<boolean>(false)
   const { theme } = useTheme()
   const tourSettings = mangoStore((s) => s.settings.tours)
-  const [isOnboarded] = useLocalStorageState(IS_ONBOARDED_KEY)
+  const [isOnBoarded] = useLocalStorageState(IS_ONBOARDED_KEY)
 
   const leverage = useMemo(() => {
     if (!group || !mangoAccount) return 0
@@ -471,7 +471,7 @@ const AccountPage = () => {
           onClose={() => setShowWithdrawModal(false)}
         />
       ) : null}
-      {!tourSettings?.account_tour_seen && isOnboarded && connected ? (
+      {!tourSettings?.account_tour_seen && isOnBoarded && connected ? (
         <AccountOnboardingTour />
       ) : null}
     </>

components/modals/UserSetupModal.tsx

@@ -53,7 +53,7 @@ const UserSetupModal = ({ isOpen, onClose }: ModalProps) => {
   const [submitDeposit, setSubmitDeposit] = useState(false)
   const [sizePercentage, setSizePercentage] = useState('')
   const [showEditProfilePic, setShowEditProfilePic] = useState(false)
-  const [, setIsOnboarded] = useLocalStorageState(IS_ONBOARDED_KEY)
+  const [, setIsOnBoarded] = useLocalStorageState(IS_ONBOARDED_KEY)
   const walletTokens = mangoStore((s) => s.wallet.tokens)
 
   const solBalance = useMemo(() => {
@@ -73,7 +73,7 @@ const UserSetupModal = ({ isOpen, onClose }: ModalProps) => {
       try {
         await handleWalletConnect(wallet)
         setShowSetupStep(2)
-        setIsOnboarded(true)
+        setIsOnBoarded(true)
       } catch (e) {
         notify({
           title: 'Setup failed. Refresh and try again.',

components/wallet/ConnectWalletButton.tsx

@@ -5,7 +5,6 @@ import { useTranslation } from 'next-i18next'
 import uniqBy from 'lodash/uniqBy'
 import WalletSelect from './WalletSelect'
 import mangoStore from '@store/mangoStore'
-import { Wallet as AnchorWallet } from '@project-serum/anchor'
 import { notify } from '../../utils/notifications'
 import Loading from '../shared/Loading'
 
@@ -13,13 +12,10 @@ export const handleWalletConnect = async (wallet: Wallet) => {
   if (!wallet) {
     return
   }
-  const actions = mangoStore.getState().actions
   const set = mangoStore.getState().set
 
   try {
     await wallet?.adapter?.connect()
-    await actions.connectMangoClientWithWallet(wallet)
-    await onConnectFetchAccountData(wallet)
     set((state) => {
       state.connected = true
     })
@@ -41,14 +37,6 @@ export const handleWalletConnect = async (wallet: Wallet) => {
   }
 }
 
-const onConnectFetchAccountData = async (wallet: Wallet) => {
-  if (!wallet) return
-  const actions = mangoStore.getState().actions
-  await actions.fetchMangoAccounts(wallet.adapter as unknown as AnchorWallet)
-  actions.fetchTourSettings(wallet.adapter.publicKey?.toString() as string)
-  actions.fetchWalletTokens(wallet.adapter as unknown as AnchorWallet)
-}
-
 export const ConnectWalletButton: React.FC = () => {
   const { wallet, wallets, select } = useWallet()
   const group = mangoStore((s) => s.group)

components/wallet/ConnectedMenu.tsx

@@ -4,7 +4,7 @@ import {
   CurrencyDollarIcon,
   UserCircleIcon,
 } from '@heroicons/react/20/solid'
-import { useWallet } from '@solana/wallet-adapter-react'
+import { useWallet, Wallet } from '@solana/wallet-adapter-react'
 import { useTranslation } from 'next-i18next'
 import { Fragment, useCallback, useEffect, useState } from 'react'
 import mangoStore from '@store/mangoStore'
@@ -16,6 +16,7 @@ import { useViewport } from 'hooks/useViewport'
 import { breakpoints } from '../../utils/theme'
 import EditProfileModal from '@components/modals/EditProfileModal'
 import MangoAccountsListModal from '@components/modals/MangoAccountsListModal'
+import { Wallet as AnchorWallet } from '@project-serum/anchor'
 
 const ConnectedMenu = () => {
   const { t } = useTranslation('common')
@@ -28,7 +29,13 @@ const ConnectedMenu = () => {
   const loadProfileDetails = mangoStore((s) => s.profile.loadDetails)
   const { width } = useViewport()
   const isMobile = width ? width < breakpoints.md : false
-
+  const onConnectFetchAccountData = async (wallet: Wallet) => {
+    if (!wallet) return
+    const actions = mangoStore.getState().actions
+    await actions.fetchMangoAccounts(wallet.adapter as unknown as AnchorWallet)
+    actions.fetchTourSettings(wallet.adapter.publicKey?.toString() as string)
+    actions.fetchWalletTokens(wallet.adapter as unknown as AnchorWallet)
+  }
   const handleDisconnect = useCallback(() => {
     set((state) => {
       state.activityFeed.feed = []
@@ -49,8 +56,14 @@ const ConnectedMenu = () => {
   }, [set, t, disconnect])
 
   useEffect(() => {
+    const handleGetWalletMangoData = async () => {
+      const actions = mangoStore.getState().actions
+      await actions.connectMangoClientWithWallet(wallet!)
+      await onConnectFetchAccountData(wallet!)
+    }
     if (publicKey) {
       actions.fetchProfileDetails(publicKey.toString())
+      handleGetWalletMangoData()
     }
   }, [publicKey])
 

store/mangoStore.ts

@@ -599,12 +599,24 @@ const mangoStore = create<MangoStore>()(
               group,
               wallet.publicKey
             )
-            if (!mangoAccounts?.length) return
+            const selectedAccountIsNotInAccountsList = mangoAccounts.find(
+              (x) =>
+                x.publicKey.toBase58() ===
+                selectedMangoAccount?.publicKey.toBase58()
+            )
+            if (!mangoAccounts?.length) {
+              set((state) => {
+                state.mangoAccounts = []
+                state.mangoAccount.current = undefined
+                state.mangoAccount.lastUpdatedAt = new Date().toISOString()
+              })
+              return
+            }
 
             mangoAccounts.forEach((ma) => ma.reloadAccountData(client))
 
             let newSelectedMangoAccount = selectedMangoAccount
-            if (!selectedMangoAccount) {
+            if (!selectedMangoAccount || !selectedAccountIsNotInAccountsList) {
               const lastAccount = localStorage.getItem(LAST_ACCOUNT_KEY)
               newSelectedMangoAccount = mangoAccounts[0]