better handle latency from wallet providers

components/ConnectWalletButton.tsx

@@ -11,15 +11,23 @@ const StyledWalletTypeLabel = styled.div`
 
 const ConnectWalletButton = () => {
   const wallet = useMangoStore((s) => s.wallet.current)
+  const set = useMangoStore((s) => s.set)
   const [savedProviderUrl] = useLocalStorageState(
     'walletProvider',
     DEFAULT_PROVIDER.url
   )
 
+  const handleWalletConect = () => {
+    wallet.connect()
+    set((state) => {
+      state.selectedMarginAccount.initialLoad = true
+    })
+  }
+
   return (
     <div className="flex justify-between border border-th-primary rounded-md h-11 w-48">
       <button
-        onClick={() => wallet.connect()}
+        onClick={handleWalletConect}
         disabled={!wallet}
         className="text-th-primary hover:text-th-fgd-1 hover:bg-th-primary focus:outline-none disabled:text-th-fgd-4 disabled:cursor-wait rounded-r-none"
       >

components/MarginBalances.tsx

@@ -22,6 +22,9 @@ export default function MarginBalances() {
   const selectedMarginAccount = useMangoStore(
     (s) => s.selectedMarginAccount.current
   )
+  const loadingMarginAccount = useMangoStore(
+    (s) => s.selectedMarginAccount.initialLoad
+  )
   const connected = useMangoStore((s) => s.wallet.connected)
   const { symbols } = useMarketList()
 
@@ -148,14 +151,16 @@ export default function MarginBalances() {
           <Button
             onClick={() => setShowDepositModal(true)}
             className="w-1/2"
-            disabled={!connected}
+            disabled={!connected || loadingMarginAccount}
           >
             <span>Deposit</span>
           </Button>
           <Button
             onClick={() => setShowWithdrawModal(true)}
             className="ml-4 w-1/2"
-            disabled={!connected || !selectedMarginAccount}
+            disabled={
+              !connected || !selectedMarginAccount || loadingMarginAccount
+            }
           >
             <span>Withdraw</span>
           </Button>

hooks/useWallet.tsx

@@ -95,9 +95,12 @@ export default function useWallet() {
     if (!wallet) return
     wallet.on('connect', async () => {
       console.log('connected wallet')
-      setMangoStore((state) => {
-        state.wallet.connected = true
-      })
+      sleep(500)
+      // wait for margin account before fetching trade history
+      await actions.fetchMarginAccounts()
+      actions.fetchWalletBalances()
+      actions.fetchMangoSrmAccounts()
+      actions.fetchTradeHistory()
       notify({
         message: 'Wallet connected',
         description:
@@ -106,12 +109,9 @@ export default function useWallet() {
           '...' +
           wallet.publicKey.toString().substr(-5),
       })
-      sleep(500)
-      // wait for margin account before fetching trade history
-      await actions.fetchMarginAccounts()
-      actions.fetchWalletBalances()
-      actions.fetchMangoSrmAccounts()
-      actions.fetchTradeHistory()
+      setMangoStore((state) => {
+        state.wallet.connected = true
+      })
     })
     wallet.on('disconnect', () => {
       console.log('disconnecting wallet')

stores/useMangoStore.tsx

@@ -94,6 +94,7 @@ interface MangoStore extends State {
   marginAccounts: MarginAccount[]
   selectedMarginAccount: {
     current: MarginAccount | null
+    initialLoad: boolean
   }
   tradeForm: {
     side: 'buy' | 'sell'
@@ -154,6 +155,7 @@ const useMangoStore = create<MangoStore>((set, get) => ({
   marginAccounts: [],
   selectedMarginAccount: {
     current: null,
+    initialLoad: false,
   },
   tradeForm: {
     side: 'buy',
@@ -224,6 +226,7 @@ const useMangoStore = create<MangoStore>((set, get) => ({
     async fetchMarginAccounts() {
       const connection = get().connection.current
       const mangoGroup = get().selectedMangoGroup.current
+      const selectedMarginAcount = get().selectedMarginAccount.current
       const wallet = get().wallet.current
       const cluster = get().connection.cluster
       const mangoClient = get().mangoClient
@@ -232,6 +235,12 @@ const useMangoStore = create<MangoStore>((set, get) => ({
 
       if (!wallet?.publicKey || !wallet.publicKey) return
 
+      if (!selectedMarginAcount) {
+        set((state) => {
+          state.selectedMarginAccount.initialLoad = true
+        })
+      }
+
       return mangoClient
         .getMarginAccountsForOwner(
           connection,
@@ -243,6 +252,7 @@ const useMangoStore = create<MangoStore>((set, get) => ({
           if (marginAccounts.length > 0) {
             set((state) => {
               state.marginAccounts = marginAccounts
+              state.selectedMarginAccount.initialLoad = false
               if (state.selectedMarginAccount.current) {
                 state.selectedMarginAccount.current = marginAccounts.find(
                   (ma) =>