fix: fetch gas token balance on approval screen (#77)

ava-labs · Oct 22, 2024 · e9bcf19 · e9bcf19
1 parent aedec29
commit e9bcf19
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 20 deletions.
diff --git a/src/contexts/BalancesProvider.tsx b/src/contexts/BalancesProvider.tsx
@@ -87,7 +87,10 @@ const BalancesContext = createContext<{
     tokenId: string
   ): Promise<void>;
   getTokenPrice(addressOrSymbol: string): number | undefined;
-  updateBalanceOnAllNetworks: (accounts: Account[]) => Promise<void>;
+  updateBalanceOnNetworks: (
+    accounts: Account[],
+    chainIds?: number[]
+  ) => Promise<void>;
   registerSubscriber: (tokenTypes: TokenType[]) => void;
   unregisterSubscriber: (tokenTypes: TokenType[]) => void;
   isTokensCached: boolean;
@@ -104,7 +107,7 @@ const BalancesContext = createContext<{
     return undefined;
   },
   async refreshNftMetadata() {}, // eslint-disable-line @typescript-eslint/no-empty-function
-  async updateBalanceOnAllNetworks() {}, // eslint-disable-line @typescript-eslint/no-empty-function
+  async updateBalanceOnNetworks() {}, // eslint-disable-line @typescript-eslint/no-empty-function
   registerSubscriber() {}, // eslint-disable-line @typescript-eslint/no-empty-function
   unregisterSubscriber() {}, // eslint-disable-line @typescript-eslint/no-empty-function
   isTokensCached: true,
@@ -268,15 +271,15 @@ export function BalancesProvider({ children }: { children: any }) {
     setIsPolling(Object.values(subscribers).some((count) => count > 0));
   }, [subscribers]);
 
-  const updateBalanceOnAllNetworks = useCallback(
-    async (accounts: Account[]) => {
-      if (!network) {
+  const updateBalanceOnNetworks = useCallback(
+    async (accounts: Account[], chainIds?: number[]) => {
+      if (!network && !chainIds?.length) {
         return;
       }
 
       const updatedBalances = await request<UpdateBalancesForNetworkHandler>({
         method: ExtensionRequest.NETWORK_BALANCES_UPDATE,
-        params: [accounts],
+        params: [accounts, chainIds],
       });
 
       dispatch({
@@ -352,7 +355,7 @@ export function BalancesProvider({ children }: { children: any }) {
         balances,
         getTokenPrice,
         refreshNftMetadata,
-        updateBalanceOnAllNetworks,
+        updateBalanceOnNetworks,
         registerSubscriber,
         unregisterSubscriber,
         isTokensCached: balances.cached ?? true,

diff --git a/src/pages/Accounts/AddWalletWithSeedPhrase.tsx b/src/pages/Accounts/AddWalletWithSeedPhrase.tsx
@@ -60,7 +60,7 @@ export function AddWalletWithSeedPhrase() {
   const { capture } = useAnalyticsContext();
   const getErrorMessage = useErrorMessage();
   const formatCurrency = useConvertedCurrencyFormatter();
-  const { updateBalanceOnAllNetworks, getTotalBalance } = useBalancesContext();
+  const { updateBalanceOnNetworks, getTotalBalance } = useBalancesContext();
 
   const { isImporting, importSeedphrase } = useImportSeedphrase();
 
@@ -96,13 +96,13 @@ export function AddWalletWithSeedPhrase() {
 
       setIsKnownPhrase(false);
       setIsBalanceLoading(true);
-      await updateBalanceOnAllNetworks(
+      await updateBalanceOnNetworks(
         addies.map((addressC) => ({ addressC })) as Account[]
       );
       setIsBalanceLoading(false);
       setAddresses(addies);
     },
-    [allAccounts, updateBalanceOnAllNetworks]
+    [allAccounts, updateBalanceOnNetworks]
   );
 
   const onContinue = useCallback(() => {

diff --git a/src/pages/Accounts/components/AccountItem.tsx b/src/pages/Accounts/components/AccountItem.tsx
@@ -61,7 +61,7 @@ export const AccountItem = forwardRef(
     const history = useHistory();
     const { capture } = useAnalyticsContext();
     const { network } = useNetworkContext();
-    const { updateBalanceOnAllNetworks } = useBalancesContext();
+    const { updateBalanceOnNetworks } = useBalancesContext();
     const [isBalanceLoading, setIsBalanceLoading] = useState(false);
 
     const isActive = isActiveAccount(account.id);
@@ -141,9 +141,9 @@ export const AccountItem = forwardRef(
 
     const getBalance = useCallback(async () => {
       setIsBalanceLoading(true);
-      await updateBalanceOnAllNetworks([account]);
+      await updateBalanceOnNetworks([account]);
       setIsBalanceLoading(false);
-    }, [account, updateBalanceOnAllNetworks]);
+    }, [account, updateBalanceOnNetworks]);
 
     return (
       <Stack

diff --git a/src/pages/ApproveAction/hooks/useFeeCustomizer.tsx b/src/pages/ApproveAction/hooks/useFeeCustomizer.tsx
@@ -19,6 +19,8 @@ import { useConnectionContext } from '@src/contexts/ConnectionProvider';
 import { UpdateActionTxDataHandler } from '@src/background/services/actions/handlers/updateTxData';
 import { ExtensionRequest } from '@src/background/connections/extensionConnection/models';
 import { useTokensWithBalances } from '@src/hooks/useTokensWithBalances';
+import { useAccountsContext } from '@src/contexts/AccountsProvider';
+import { useBalancesContext } from '@src/contexts/BalancesProvider';
 
 const getInitialFeeRate = (data?: SigningData): bigint | undefined => {
   if (data?.type === RpcMethod.BITCOIN_SEND_TRANSACTION) {
@@ -38,6 +40,10 @@ export const useFeeCustomizer = ({
   network?: NetworkWithCaipId;
 }) => {
   const { action } = useApproveAction<DisplayData>(actionId);
+  const {
+    accounts: { active: activeAccount },
+  } = useAccountsContext();
+  const { updateBalanceOnNetworks } = useBalancesContext();
   const { request } = useConnectionContext();
   const [networkFee, setNetworkFee] = useState<NetworkFee | null>();
 
@@ -133,6 +139,15 @@ export const useFeeCustomizer = ({
     return nativeToken.balance > need;
   }, [getFeeInfo, nativeToken?.balance, signingData]);
 
+  // Make sure we have gas token balances for the transaction's chain
+  useEffect(() => {
+    if (!activeAccount || !network?.chainId) {
+      return;
+    }
+
+    updateBalanceOnNetworks([activeAccount], [network.chainId]);
+  }, [activeAccount, network?.chainId, updateBalanceOnNetworks]);
+
   useEffect(() => {
     const nativeBalance = nativeToken?.balance;
 

diff --git a/src/pages/ImportPrivateKey/ImportPrivateKey.tsx b/src/pages/ImportPrivateKey/ImportPrivateKey.tsx
@@ -36,7 +36,7 @@ type DerivedAddresses = {
 
 export function ImportPrivateKey() {
   const { currency, currencyFormatter } = useSettingsContext();
-  const { updateBalanceOnAllNetworks } = useBalancesContext();
+  const { updateBalanceOnNetworks } = useBalancesContext();
   const { network } = useNetworkContext();
   const { capture } = useAnalyticsContext();
   const theme = useTheme();
@@ -97,13 +97,13 @@ export function ImportPrivateKey() {
   }, [network?.isTestnet, privateKey]);
 
   useEffect(() => {
-    if (derivedAddresses && updateBalanceOnAllNetworks) {
+    if (derivedAddresses && updateBalanceOnNetworks) {
       setIsBalanceLoading(true);
-      updateBalanceOnAllNetworks([derivedAddresses as Account]).finally(() =>
+      updateBalanceOnNetworks([derivedAddresses as Account]).finally(() =>
         setIsBalanceLoading(false)
       );
     }
-  }, [derivedAddresses, updateBalanceOnAllNetworks]);
+  }, [derivedAddresses, updateBalanceOnNetworks]);
 
   return (
     <Stack

diff --git a/src/pages/Permissions/components/AccountsDropdown.tsx b/src/pages/Permissions/components/AccountsDropdown.tsx
@@ -45,7 +45,7 @@ export const AccountsDropdown = ({
   const { t } = useTranslation();
   const [selectedAccount, setSelectedAccount] = useState(activeAccount);
   const [isBalanceLoading, setIsBalanceLoading] = useState(false);
-  const { updateBalanceOnAllNetworks } = useBalancesContext();
+  const { updateBalanceOnNetworks } = useBalancesContext();
   const { currency, currencyFormatter } = useSettingsContext();
   const accountBalance = useBalanceTotalInCurrency(selectedAccount);
 
@@ -57,7 +57,7 @@ export const AccountsDropdown = ({
   useEffect(() => {
     const getBalance = async () => {
       setIsBalanceLoading(true);
-      await updateBalanceOnAllNetworks?.(selectedAccount);
+      await updateBalanceOnNetworks?.(selectedAccount);
       setIsBalanceLoading(false);
     };
 
@@ -69,7 +69,7 @@ export const AccountsDropdown = ({
     }
 
     getBalance();
-  }, [activeAccount, selectedAccount, updateBalanceOnAllNetworks]);
+  }, [activeAccount, selectedAccount, updateBalanceOnNetworks]);
 
   // Update balance & notify parent component about changes when account is selected
   useEffect(() => {