Fix account balance discrepancies

src/__swaps__/screens/Swap/providers/swap-provider.tsx

@@ -264,24 +264,20 @@ export const SwapProvider = ({ children }: SwapProviderProps) => {
         }
       }
 
-      queryClient.invalidateQueries([
-        // old user assets invalidation (will cause a re-fetch)
-        {
-          queryKey: userAssetsQueryKey({
-            address: parameters.quote.from,
-            currency: nativeCurrency,
-            connectedToHardhat,
-          }),
-        },
-        // new swaps user assets invalidations
-        {
-          queryKey: swapsUserAssetsQueryKey({
-            address: parameters.quote.from as Address,
-            currency: nativeCurrency,
-            testnetMode: !!connectedToHardhat,
-          }),
-        },
-      ]);
+      queryClient.invalidateQueries(
+        userAssetsQueryKey({
+          address: parameters.quote.from,
+          currency: nativeCurrency,
+          connectedToHardhat,
+        })
+      );
+      queryClient.invalidateQueries(
+        swapsUserAssetsQueryKey({
+          address: parameters.quote.from as Address,
+          currency: nativeCurrency,
+          testnetMode: !!connectedToHardhat,
+        })
+      );
 
       clearCustomGasSettings(chainId);
       NotificationManager?.postNotification('rapCompleted');

src/components/DappBrowser/control-panel/ControlPanel.tsx

@@ -158,7 +158,7 @@ export const ControlPanel = () => {
               wallet.type === WalletTypes.readOnly
                 ? i18n.t(i18n.l.wallet.change_wallet.watching)
                 : account.balance
-                  ? convertAmountToNativeDisplay(account.balance, nativeCurrency)
+                  ? account.balance
                   : i18n.t(i18n.l.wallet.change_wallet.no_balance),
             uniqueId: account.address,
             color: colors.avatarBackgrounds[account.color],
@@ -472,7 +472,7 @@ const HomePanel = ({
       },
       screen: Routes.MAIN_EXCHANGE_SCREEN,
     });
-  }, [runWalletChecksBeforeSwapOrBridge, selectedWallet?.uniqueId]);
+  }, [navigate, runWalletChecksBeforeSwapOrBridge, selectedWallet?.uniqueId, swapsV2Enabled]);
 
   const handleOnPressBridge = useCallback(async () => {
     const valid = await runWalletChecksBeforeSwapOrBridge();

src/components/change-wallet/AddressRow.tsx

@@ -22,9 +22,6 @@ import { EditWalletContextMenuActions } from '@/screens/ChangeWalletSheet';
 import { toChecksumAddress } from '@/handlers/web3';
 import { IS_IOS, IS_ANDROID } from '@/env';
 import { ContextMenu } from '../context-menu';
-import { convertAmountToNativeDisplay } from '@/helpers/utilities';
-import { useSelector } from 'react-redux';
-import { AppState } from '@/redux/store';
 import { useForegroundColor } from '@/design-system';
 
 const maxAccountLabelWidth = deviceUtils.dimensions.width - 88;
@@ -121,7 +118,6 @@ interface AddressRowProps {
 }
 
 export default function AddressRow({ contextMenuActions, data, editMode, onPress }: AddressRowProps) {
-  const nativeCurrency = useSelector((state: AppState) => state.settings.nativeCurrency);
   const notificationsEnabled = useExperimentalFlag(NOTIFICATIONS);
 
   const { address, balance, color: accountColor, ens, image: accountImage, isSelected, isReadOnly, isLedger, label, walletId } = data;
@@ -130,13 +126,13 @@ export default function AddressRow({ contextMenuActions, data, editMode, onPress
 
   const labelQuaternary = useForegroundColor('labelQuaternary');
 
-  const cleanedUpBalance = useMemo(() => {
+  const balanceOrNoBalance = useMemo(() => {
     if (balance) {
-      return convertAmountToNativeDisplay(balance, nativeCurrency);
+      return balance;
     } else {
       return lang.t('wallet.change_wallet.no_balance');
     }
-  }, [balance, nativeCurrency]);
+  }, [balance]);
 
   const cleanedUpLabel = useMemo(() => removeFirstEmojiFromString(label), [label]);
 
@@ -254,7 +250,7 @@ export default function AddressRow({ contextMenuActions, data, editMode, onPress
               <StyledTruncatedText color={colors.dark} testID={`change-wallet-address-row-label-${walletName}`}>
                 {walletName}
               </StyledTruncatedText>
-              <StyledBottomRowText color={labelQuaternary}>{cleanedUpBalance}</StyledBottomRowText>
+              <StyledBottomRowText color={labelQuaternary}>{balanceOrNoBalance}</StyledBottomRowText>
             </ColumnWithMargins>
           </Row>
           <Column style={sx.rightContent}>

src/components/contacts/ContactRow.js

@@ -11,12 +11,11 @@ import useExperimentalFlag, { PROFILES } from '@/config/experimentalHooks';
 import { fetchReverseRecord } from '@/handlers/ens';
 import { ENS_DOMAIN } from '@/helpers/ens';
 import { isENSAddressFormat, isValidDomainFormat } from '@/helpers/validators';
-import { useAccountSettings, useContacts, useDimensions, useENSAvatar } from '@/hooks';
+import { useContacts, useDimensions, useENSAvatar } from '@/hooks';
 import styled from '@/styled-thing';
 import { margin } from '@/styles';
 import { addressHashedColorIndex, addressHashedEmoji } from '@/utils/profileUtils';
 import * as i18n from '@/languages';
-import { convertAmountToNativeDisplay } from '@/helpers/utilities';
 
 const ContactAddress = styled(TruncatedAddress).attrs(({ theme: { colors }, lite }) => ({
   align: 'left',
@@ -58,17 +57,16 @@ const ContactRow = ({ address, color, nickname, symmetricalMargins, ...props },
   const profilesEnabled = useExperimentalFlag(PROFILES);
   const { width: deviceWidth } = useDimensions();
   const { onAddOrUpdateContacts } = useContacts();
-  const { nativeCurrency } = useAccountSettings();
   const { colors } = useTheme();
   const { accountType, balance, ens, image, label, network, onPress, showcaseItem, testID } = props;
 
   const cleanedUpBalance = useMemo(() => {
     if (balance) {
-      return convertAmountToNativeDisplay(balance, nativeCurrency);
+      return balance;
     } else {
       return i18n.t(i18n.l.wallet.change_wallet.no_balance);
     }
-  }, [balance, nativeCurrency]);
+  }, [balance]);
 
   // show avatar for contact rows that are accounts, not contacts
   const avatar = accountType !== 'contacts' ? returnStringFirstEmoji(label) || profileUtils.addressHashedEmoji(address) : null;

src/helpers/buildWalletSections.tsx

@@ -1,6 +1,5 @@
 import { createSelector } from 'reselect';
 import { buildBriefCoinsList, buildBriefUniqueTokenList } from './assets';
-import { add, convertAmountToNativeDisplay } from './utilities';
 import { NativeCurrencyKey, ParsedAddressAsset } from '@/entities';
 import { queryClient } from '@/react-query';
 import { positionsQueryKey } from '@/resources/defi/PositionsQuery';
@@ -39,6 +38,7 @@ const EMPTY_WALLET_CONTENT = [
 const ONLY_NFTS_CONTENT = [{ type: 'ETH_CARD', uid: 'eth-card' }];
 
 const sortedAssetsSelector = (state: any) => state.sortedAssets;
+const accountBalanceSelector = (state: any) => state.accountBalance;
 const hiddenCoinsSelector = (state: any) => state.hiddenCoins;
 const isCoinListEditedSelector = (state: any) => state.isCoinListEdited;
 const isLoadingUserAssetsSelector = (state: any) => state.isLoadingUserAssets;
@@ -105,23 +105,14 @@ const withPositionsSection = (isLoadingUserAssets: boolean) => {
 const withBriefBalanceSection = (
   sortedAssets: ParsedAddressAsset[],
   isLoadingUserAssets: boolean,
+  accountBalance: number | undefined,
   nativeCurrency: NativeCurrencyKey,
   isCoinListEdited: boolean,
   pinnedCoins: any,
   hiddenCoins: any,
-  collectibles: any,
-  nftSort: string
+  collectibles: any
 ) => {
-  const { briefAssets, totalBalancesValue } = buildBriefCoinsList(sortedAssets, nativeCurrency, isCoinListEdited, pinnedCoins, hiddenCoins);
-
-  const { accountAddress: address } = store.getState().settings;
-  const positionsObj: RainbowPositions | undefined = queryClient.getQueryData(positionsQueryKey({ address, currency: nativeCurrency }));
-
-  const positionsTotal = positionsObj?.totals?.total?.amount || '0';
-
-  const totalBalanceWithPositionsValue = add(totalBalancesValue, positionsTotal);
-
-  const totalValue = convertAmountToNativeDisplay(totalBalanceWithPositionsValue, nativeCurrency);
+  const { briefAssets } = buildBriefCoinsList(sortedAssets, nativeCurrency, isCoinListEdited, pinnedCoins, hiddenCoins);
 
   const hasTokens = briefAssets?.length;
   const hasNFTs = collectibles?.length;
@@ -133,7 +124,7 @@ const withBriefBalanceSection = (
     {
       type: 'PROFILE_STICKY_HEADER',
       uid: 'assets-profile-header-compact',
-      value: totalValue,
+      value: accountBalance,
       isLoadingUserAssets,
     },
     {
@@ -162,7 +153,7 @@ const withBriefBalanceSection = (
           {
             type: 'PROFILE_BALANCE_ROW',
             uid: 'profile-balance',
-            value: totalValue,
+            value: accountBalance,
           },
           {
             type: 'PROFILE_BALANCE_ROW_SPACE_AFTER',
@@ -172,13 +163,13 @@ const withBriefBalanceSection = (
     {
       type: 'PROFILE_ACTION_BUTTONS_ROW',
       uid: 'profile-action-buttons',
-      value: totalValue,
+      value: accountBalance,
     },
     hasTokens
       ? {
           type: 'PROFILE_ACTION_BUTTONS_ROW_SPACE_AFTER',
           uid: 'profile-action-buttons-space-after',
-          value: totalValue,
+          value: accountBalance,
         }
       : { type: 'BIG_EMPTY_WALLET_SPACER', uid: 'big-empty-wallet-spacer-1' },
   ];
@@ -226,6 +217,7 @@ const briefBalanceSectionSelector = createSelector(
   [
     sortedAssetsSelector,
     isLoadingUserAssetsSelector,
+    accountBalanceSelector,
     nativeCurrencySelector,
     isCoinListEditedSelector,
     pinnedCoinsSelector,

src/hooks/useRefreshAccountData.ts

@@ -1,6 +1,6 @@
 import { captureException } from '@sentry/react-native';
 import delay from 'delay';
-import { useCallback, useState } from 'react';
+import { useCallback, useMemo, useState } from 'react';
 import { useDispatch } from 'react-redux';
 import { getIsHardhatConnected } from '@/handlers/web3';
 import { walletConnectLoadState } from '../redux/walletconnect';
@@ -10,34 +10,36 @@ import { PROFILES, useExperimentalFlag } from '@/config';
 import logger from '@/utils/logger';
 import { queryClient } from '@/react-query';
 import { userAssetsQueryKey } from '@/resources/assets/UserAssetsQuery';
+import { userAssetsQueryKey as swapsUserAssetsQueryKey } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { nftsQueryKey } from '@/resources/nfts';
 import { positionsQueryKey } from '@/resources/defi/PositionsQuery';
+import { Address } from 'viem';
+import { addysSummaryQueryKey } from '@/resources/summary/summary';
+import useWallets from './useWallets';
 
 export default function useRefreshAccountData() {
   const dispatch = useDispatch();
   const { accountAddress, nativeCurrency } = useAccountSettings();
   const [isRefreshing, setIsRefreshing] = useState(false);
   const profilesEnabled = useExperimentalFlag(PROFILES);
 
+  const { wallets } = useWallets();
+
+  const allAddresses = useMemo(
+    () => Object.values(wallets || {}).flatMap(wallet => wallet.addresses.map(account => account.address as Address)),
+    [wallets]
+  );
+
   const fetchAccountData = useCallback(async () => {
     const connectedToHardhat = getIsHardhatConnected();
 
-    queryClient.invalidateQueries({
-      queryKey: nftsQueryKey({ address: accountAddress }),
-    });
-    queryClient.invalidateQueries({
-      queryKey: positionsQueryKey({
-        address: accountAddress,
-        currency: nativeCurrency,
-      }),
-    });
-    queryClient.invalidateQueries({
-      queryKey: userAssetsQueryKey({
-        address: accountAddress,
-        currency: nativeCurrency,
-        connectedToHardhat,
-      }),
-    });
+    queryClient.invalidateQueries(nftsQueryKey({ address: accountAddress }));
+    queryClient.invalidateQueries(positionsQueryKey({ address: accountAddress as Address, currency: nativeCurrency }));
+    queryClient.invalidateQueries(addysSummaryQueryKey({ addresses: allAddresses, currency: nativeCurrency }));
+    queryClient.invalidateQueries(userAssetsQueryKey({ address: accountAddress, currency: nativeCurrency, connectedToHardhat }));
+    queryClient.invalidateQueries(
+      swapsUserAssetsQueryKey({ address: accountAddress as Address, currency: nativeCurrency, testnetMode: !!connectedToHardhat })
+    );
 
     try {
       const getWalletNames = dispatch(fetchWalletNames());

src/hooks/useWalletBalances.ts

@@ -3,19 +3,88 @@ import { useMemo } from 'react';
 import { Address } from 'viem';
 import useAccountSettings from './useAccountSettings';
 import { useAddysSummary } from '@/resources/summary/summary';
+import { useQueries } from '@tanstack/react-query';
+import { fetchPositions, positionsQueryKey } from '@/resources/defi/PositionsQuery';
+import { RainbowPositions } from '@/resources/defi/types';
+import { add, convertAmountToNativeDisplay } from '@/helpers/utilities';
+import { queryClient } from '@/react-query';
 
-const useWalletBalances = (wallets: AllRainbowWallets) => {
+type WalletBalanceWithPositions = {
+  assetBalanceAmount: number;
+  assetBalanceDisplay: string;
+  positionsBalanceAmount: string | number;
+  positionsBalanceDisplay: string;
+  totalBalanceAmount: string | number;
+  totalBalanceDisplay: string;
+};
+
+type WalletBalanceWithoutPositions = Omit<WalletBalanceWithPositions, 'positionsBalanceAmount' | 'positionsBalanceDisplay'>;
+
+type WalletBalanceResult<T extends boolean> = {
+  balances: Record<Address, T extends true ? WalletBalanceWithPositions : WalletBalanceWithoutPositions>;
+  isLoading: boolean;
+};
+
+const useWalletBalances = <T extends boolean = true>(wallets: AllRainbowWallets, withPositions: T = true as T): WalletBalanceResult<T> => {
   const { nativeCurrency } = useAccountSettings();
 
-  const walletAddresses: Address[] = useMemo(
+  const allAddresses = useMemo(
     () => Object.values(wallets).flatMap(wallet => wallet.addresses.map(account => account.address as Address)),
     [wallets]
   );
 
-  const { data, isLoading } = useAddysSummary({ addresses: walletAddresses, currency: nativeCurrency });
+  const { data: summaryData, isLoading: isSummaryLoading } = useAddysSummary({
+    addresses: allAddresses,
+    currency: nativeCurrency,
+  });
+
+  const positionQueries = useQueries({
+    queries: allAddresses.map(address => ({
+      queryKey: positionsQueryKey({ address, currency: nativeCurrency }),
+      queryFn: () => fetchPositions({ address, currency: nativeCurrency }),
+      enabled: !!address && (withPositions as boolean),
+    })),
+  });
+
+  const balancesWithPositions = useMemo(() => {
+    const result: Record<Address, WalletBalanceWithPositions | WalletBalanceWithoutPositions> = {};
+
+    for (const address of allAddresses) {
+      const lowerCaseAddress = address.toLowerCase() as Address;
+      const assetBalance = summaryData?.data?.addresses?.[lowerCaseAddress]?.summary?.asset_value || 0;
+
+      if (!withPositions) {
+        result[lowerCaseAddress] = {
+          assetBalanceAmount: assetBalance,
+          assetBalanceDisplay: convertAmountToNativeDisplay(assetBalance, nativeCurrency),
+          totalBalanceAmount: assetBalance,
+          totalBalanceDisplay: convertAmountToNativeDisplay(assetBalance, nativeCurrency),
+        };
+      } else {
+        const positionData = queryClient.getQueryData<RainbowPositions | undefined>(
+          positionsQueryKey({ address, currency: nativeCurrency })
+        );
+        const positionsBalance = positionData?.totals?.total?.amount || '0';
+        const totalAccountBalance = add(assetBalance, positionsBalance);
+
+        result[lowerCaseAddress] = {
+          assetBalanceAmount: assetBalance,
+          assetBalanceDisplay: convertAmountToNativeDisplay(assetBalance, nativeCurrency),
+          positionsBalanceAmount: positionsBalance,
+          positionsBalanceDisplay: convertAmountToNativeDisplay(positionsBalance, nativeCurrency),
+          totalBalanceAmount: totalAccountBalance,
+          totalBalanceDisplay: convertAmountToNativeDisplay(totalAccountBalance, nativeCurrency),
+        };
+      }
+    }
+
+    return result;
+  }, [allAddresses, summaryData, withPositions, nativeCurrency]);
+
+  const isLoading = isSummaryLoading || positionQueries.some(query => query.isLoading);
 
   return {
-    balances: data?.data?.addresses,
+    balances: balancesWithPositions as WalletBalanceResult<T>['balances'],
     isLoading,
   };
 };