From 605637a9dee1fce932e4a3cdb8dbe9353c58ec1f Mon Sep 17 00:00:00 2001
From: saml33 <slam.uke@gmail.com>
Date: Fri, 29 Sep 2023 11:53:12 +1000
Subject: [PATCH] refetch wallet balance on wallet swap

---
 components/swap/SwapForm.tsx            | 6 +++---
 components/swap/SwapReviewRouteInfo.tsx | 2 ++
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/components/swap/SwapForm.tsx b/components/swap/SwapForm.tsx
index 3cfa9586..135ef7fd 100644
--- a/components/swap/SwapForm.tsx
+++ b/components/swap/SwapForm.tsx
@@ -31,7 +31,7 @@ const set = mangoStore.getState().set
 const SwapForm = () => {
   const { t } = useTranslation(['common', 'swap', 'trade'])
   const groupLoaded = mangoStore((s) => s.groupLoaded)
-  const { mangoAccountAddress } = useMangoAccount()
+  const { mangoAccountAddress, initialLoad } = useMangoAccount()
   const { connected } = useWallet()
   const { query } = useRouter()
   const [showTokenSelect, setShowTokenSelect] =
@@ -55,10 +55,10 @@ const SwapForm = () => {
 
   // enable wallet swap when connected and no mango account
   useEffect(() => {
-    if (connected && !mangoAccountAddress) {
+    if (connected && !mangoAccountAddress && !initialLoad) {
       setWalletSwap(true)
     }
-  }, [connected, mangoAccountAddress])
+  }, [connected, mangoAccountAddress, initialLoad])
 
   // setup swap from url query
   useEffect(() => {
diff --git a/components/swap/SwapReviewRouteInfo.tsx b/components/swap/SwapReviewRouteInfo.tsx
index b4541bea..d09a93d5 100644
--- a/components/swap/SwapReviewRouteInfo.tsx
+++ b/components/swap/SwapReviewRouteInfo.tsx
@@ -256,6 +256,7 @@ const SwapReviewRouteInfo = ({
 
   const onWalletSwap = useCallback(async () => {
     if (!selectedRoute || !inputBank || !outputBank || !wallet.publicKey) return
+    const actions = mangoStore.getState().actions
     const client = mangoStore.getState().client
     const connection = mangoStore.getState().connection
     setSubmitting(true)
@@ -276,6 +277,7 @@ const SwapReviewRouteInfo = ({
         type: 'success',
         txid: tx.signature,
       })
+      actions.fetchWalletTokens(wallet.publicKey)
     } catch (e) {
       console.log('error swapping wallet tokens', e)
     } finally {