fix: improve token balance and transaction history invalidation

lifinance · Nov 6, 2024 · 46fc046 · 46fc046
1 parent 28a80af
commit 46fc046
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 52 deletions.
diff --git a/packages/widget/src/hooks/useGasRefuel.ts b/packages/widget/src/hooks/useGasRefuel.ts
@@ -19,23 +19,30 @@ export const useGasRefuel = () => {
   const toChain = getChainById(toChainId)
   const fromChain = getChainById(fromChainId)
 
-  const { account } = useAccount({ chainType: fromChain?.chainType })
+  const { accounts } = useAccount()
+
+  const fromAccount = accounts.find(
+    (account) => account.chainType === fromChain?.chainType
+  )
+
+  const toAccount = accounts.find(
+    (account) => account.chainType === toChain?.chainType
+  )
 
   const isFromContractAddress = useIsContractAddress(
-    account.address,
+    fromAccount?.address,
     fromChainId,
-    account.chainType
+    fromAccount?.chainType
   )
   const isToContractAddress = useIsContractAddress(
     toAddress,
     toChainId,
     toChain?.chainType
   )
 
-  const { token: nativeToken } = useTokenBalance(
-    toAddress,
-    toChainId ? toChain?.nativeToken : undefined,
-    toChain
+  const { token: destinationNativeToken } = useTokenBalance(
+    toAddress || toAccount?.address,
+    toChainId ? toChain?.nativeToken : undefined
   )
 
   const { data: gasRecommendation, isLoading } = useGasRecommendation(
@@ -49,7 +56,7 @@ export const useGasRefuel = () => {
     fromChain?.chainType !== toChain?.chainType ? Boolean(toAddress) : true
 
   const isToAddressSatisfied = isFromContractAddress
-    ? toAddress && toAddress !== account.address && !isToContractAddress
+    ? toAddress && toAddress !== fromAccount?.address && !isToContractAddress
     : true
 
   const enabled = useMemo(() => {
@@ -59,13 +66,13 @@ export const useGasRefuel = () => {
       fromChainId === toChainId ||
       !gasRecommendation?.available ||
       !gasRecommendation?.recommended ||
-      !nativeToken ||
+      !destinationNativeToken ||
       !isChainTypeSatisfied ||
       !isToAddressSatisfied
     ) {
       return false
     }
-    const tokenBalance = nativeToken.amount ?? 0n
+    const tokenBalance = destinationNativeToken.amount ?? 0n
 
     // Check if the user balance < 50% of the recommended amount
     const recommendedAmount = BigInt(gasRecommendation.recommended.amount) / 2n
@@ -77,7 +84,7 @@ export const useGasRefuel = () => {
     gasRecommendation,
     isChainTypeSatisfied,
     isToAddressSatisfied,
-    nativeToken,
+    destinationNativeToken,
     toChainId,
   ])
 

diff --git a/packages/widget/src/hooks/useRouteExecution.ts b/packages/widget/src/hooks/useRouteExecution.ts
@@ -59,16 +59,42 @@ export const useRouteExecution = ({
         process,
       })
     }
-
-    if (isRouteDone(clonedUpdatedRoute)) {
+    const executionCompleted = isRouteDone(clonedUpdatedRoute)
+    const executionFailed = isRouteFailed(clonedUpdatedRoute)
+    if (executionCompleted) {
       emitter.emit(WidgetEvent.RouteExecutionCompleted, clonedUpdatedRoute)
     }
-    if (isRouteFailed(clonedUpdatedRoute) && process) {
+    if (executionFailed && process) {
       emitter.emit(WidgetEvent.RouteExecutionFailed, {
         route: clonedUpdatedRoute,
         process,
       })
     }
+    if (executionCompleted || executionFailed) {
+      const invalidateKeys = [
+        [
+          'token-balances',
+          clonedUpdatedRoute.fromAddress,
+          clonedUpdatedRoute.fromChainId,
+        ],
+        [
+          'token-balances',
+          clonedUpdatedRoute.toAddress,
+          clonedUpdatedRoute.toChainId,
+        ],
+        ['transaction-history'],
+      ]
+      for (const key of invalidateKeys) {
+        queryClient.invalidateQueries(
+          {
+            queryKey: key,
+            exact: false,
+            refetchType: 'all',
+          },
+          { cancelRefetch: false }
+        )
+      }
+    }
     // biome-ignore lint/suspicious/noConsole: logs route information
     console.log('Route updated.', clonedUpdatedRoute)
   }

diff --git a/packages/widget/src/hooks/useTokenBalance.ts b/packages/widget/src/hooks/useTokenBalance.ts
@@ -1,27 +1,21 @@
-import type { ExtendedChain } from '@lifi/sdk'
 import { type Token, type TokenAmount, getTokenBalances } from '@lifi/sdk'
-import { useAccount } from '@lifi/wallet-management'
 import { useQuery, useQueryClient } from '@tanstack/react-query'
 import { useCallback, useMemo } from 'react'
 
 const defaultRefetchInterval = 30_000
 
-export const useTokenBalance = (
-  accountAddress?: string,
-  token?: Token,
-  chain?: ExtendedChain
-) => {
-  const { account } = useAccount(
-    // When we provide chain we want to be sure that account address used is from the same ecosystem as token
-    chain ? { chainType: chain.chainType } : undefined
-  )
+export const useTokenBalance = (accountAddress?: string, token?: Token) => {
   const queryClient = useQueryClient()
-  const walletAddress = accountAddress || account.address
 
   const tokenBalanceQueryKey = useMemo(
     () =>
-      ['token-balance', walletAddress, token?.chainId, token?.address] as const,
-    [token?.address, token?.chainId, walletAddress]
+      [
+        'token-balance',
+        accountAddress,
+        token?.chainId,
+        token?.address,
+      ] as const,
+    [token?.address, token?.chainId, accountAddress]
   )
 
   const { data, isLoading, refetch } = useQuery({
@@ -73,18 +67,11 @@ export const useTokenBalance = (
       } as TokenAmount
     },
 
-    enabled: Boolean(walletAddress && token),
+    enabled: Boolean(accountAddress && token),
     refetchInterval: defaultRefetchInterval,
     staleTime: defaultRefetchInterval,
   })
 
-  const refetchAllBalances = () => {
-    queryClient.refetchQueries({
-      queryKey: ['token-balances', accountAddress, token?.chainId],
-      exact: false,
-    })
-  }
-
   const refetchNewBalance = useCallback(() => {
     queryClient.setQueryDefaults(tokenBalanceQueryKey, {
       refetchInterval: 250,
@@ -97,7 +84,6 @@ export const useTokenBalance = (
     isLoading,
     refetch,
     refetchNewBalance,
-    refetchAllBalances,
     getTokenBalancesWithRetry,
   }
 }

diff --git a/packages/widget/src/pages/TransactionPage/StatusBottomSheet.tsx b/packages/widget/src/pages/TransactionPage/StatusBottomSheet.tsx
@@ -5,7 +5,6 @@ import {
   WarningRounded,
 } from '@mui/icons-material'
 import { Box, Button, Skeleton, Typography } from '@mui/material'
-import { useQueryClient } from '@tanstack/react-query'
 import { useCallback, useEffect, useRef } from 'react'
 import { useTranslation } from 'react-i18next'
 import { BottomSheet } from '../../components/BottomSheet/BottomSheet.js'
@@ -78,7 +77,6 @@ export const StatusBottomSheetContent: React.FC<
 > = ({ status, route, onClose }) => {
   const { t } = useTranslation()
   const { navigateBack, navigate } = useNavigateBack()
-  const queryClient = useQueryClient()
   const { setFieldValue } = useFieldActions()
   const {
     subvariant,
@@ -100,25 +98,22 @@ export const StatusBottomSheetContent: React.FC<
     ),
   }
 
-  const toChain = getChainById(toToken.chainId)
-
-  const { token, refetch, refetchNewBalance, refetchAllBalances } =
-    useTokenBalance(route.toAddress, toToken, toChain)
+  const { token, refetch, refetchNewBalance } = useTokenBalance(
+    route.toAddress,
+    toToken
+  )
 
-  const invalidateQueries = () => {
-    refetchAllBalances()
+  const cleanFields = () => {
     setFieldValue('fromAmount', '')
     setFieldValue('toAmount', '')
-    queryClient.invalidateQueries({ queryKey: ['transaction-history'] })
   }
 
   const handleDone = () => {
-    invalidateQueries()
+    cleanFields()
     navigateBack()
   }
 
   const handlePartialDone = () => {
-    invalidateQueries()
     if (
       toToken.chainId !== route.toToken.chainId &&
       toToken.address !== route.toToken.address
@@ -136,12 +131,14 @@ export const StatusBottomSheetContent: React.FC<
       setFieldValue('toToken', route.toToken.address, {
         isTouched: true,
       })
+    } else {
+      cleanFields()
     }
     navigateBack()
   }
 
   const handleClose = () => {
-    invalidateQueries()
+    cleanFields()
     onClose()
   }
 
@@ -178,7 +175,7 @@ export const StatusBottomSheetContent: React.FC<
         primaryMessage = t('success.message.exchangeSuccessful', {
           amount: formatTokenAmount(token.amount, token.decimals),
           tokenSymbol: token.symbol,
-          chainName: toChain?.name,
+          chainName: getChainById(toToken.chainId)?.name,
           walletAddress: shortenAddress(route.toAddress),
         })
       }
@@ -195,7 +192,7 @@ export const StatusBottomSheetContent: React.FC<
         secondaryMessage = t('success.message.exchangeSuccessful', {
           amount: formatTokenAmount(token.amount, token.decimals),
           tokenSymbol: token.symbol,
-          chainName: toChain?.name,
+          chainName: getChainById(toToken.chainId)?.name,
           walletAddress: shortenAddress(route.toAddress),
         })
       }
@@ -212,7 +209,7 @@ export const StatusBottomSheetContent: React.FC<
         secondaryMessage = t('success.message.exchangeSuccessful', {
           amount: formatTokenAmount(token.amount, token.decimals),
           tokenSymbol: token.symbol,
-          chainName: toChain?.name,
+          chainName: getChainById(toToken.chainId)?.name,
           walletAddress: shortenAddress(route.toAddress),
         })
       }