@rango-dev/widget-embedded 0.41.1-next.4 → 0.41.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rango-dev/widget-embedded",
-  "version": "0.41.1-next.4",
+  "version": "0.41.1",
   "license": "MIT",
   "type": "module",
   "source": "./src/index.ts",

package/src/hooks/useSubscribeToWidgetEvents/useSubscribeToWidgetEvents.ts

@@ -1,4 +1,5 @@
-import type { RouteEventData, StepEventData } from '../..';
+import type { ConnectedWallet, RouteEventData, StepEventData } from '../..';
+import type { Asset } from 'rango-types';
 
 import {
   isApprovalTX,
@@ -15,11 +16,21 @@ import { useNotificationStore } from '../../store/notification';
 
 export function useSubscribeToWidgetEvents() {
   const setNotification = useNotificationStore.use.setNotification();
-  const { connectedWallets, fetchBalances } = useAppStore();
+  const {
+    connectedWallets,
+    fetchBalances,
+    customTokens: getCustomTokens,
+  } = useAppStore();
 
   useEffect(() => {
     const handleStepEvent = (widgetEvent: StepEventData) => {
       const { event, step, route } = widgetEvent;
+
+      /**
+       * Determine if we should refetch balances:
+       * - When a transaction is sent (TX_SENT) and it's not an approval transaction.
+       * - When a step succeeds.
+       */
       const shouldRefetchBalance =
         (event.type === StepEventType.TX_EXECUTION &&
           event.status === StepExecutionEventStatus.TX_SENT &&
@@ -27,6 +38,12 @@ export function useSubscribeToWidgetEvents() {
         event.type === StepEventType.SUCCEEDED;
 
       if (shouldRefetchBalance) {
+        const fetchBalanceAccounts: ConnectedWallet[] = [];
+
+        /**
+         * Identify the wallet corresponding to the `from` blockchain and match it with a connected wallet.
+         * If found, add it to the balance fetch list.
+         */
         const fromWallet = route.wallets[step?.fromBlockchain];
         if (fromWallet) {
           const fromAccount = connectedWallets.find(
@@ -37,10 +54,14 @@ export function useSubscribeToWidgetEvents() {
               connectedWallet.chain === step?.fromBlockchain
           );
           if (fromAccount) {
-            void fetchBalances([fromAccount]);
+            fetchBalanceAccounts.push(fromAccount);
           }
         }
 
+        /**
+         * If the transaction is cross-chain (`fromBlockchain !== toBlockchain`),
+         * find and add the `to` blockchain wallet to the balance fetch list.
+         */
         if (step?.fromBlockchain !== step?.toBlockchain) {
           const toWallet = route.wallets[step?.toBlockchain];
           if (toWallet) {
@@ -52,10 +73,41 @@ export function useSubscribeToWidgetEvents() {
                 connectedWallet.chain === step?.toBlockchain
             );
             if (toAccount) {
-              void fetchBalances([toAccount]);
+              fetchBalanceAccounts.push(toAccount);
             }
           }
         }
+
+        if (fetchBalanceAccounts.length > 0) {
+          const customTokens = getCustomTokens();
+          const stepAssets: Asset[] = [
+            {
+              blockchain: step.fromBlockchain,
+              address: step.fromSymbolAddress,
+              symbol: step.fromSymbol,
+            },
+            {
+              blockchain: step.toBlockchain,
+              address: step.toSymbolAddress,
+              symbol: step.toSymbol,
+            },
+          ];
+          const selectedCustomTokens = stepAssets.filter((asset) =>
+            customTokens.some(
+              (token) =>
+                token.blockchain === asset.blockchain &&
+                token.address?.toLocaleLowerCase() ===
+                  asset.address?.toLocaleLowerCase() &&
+                token.symbol?.toLocaleLowerCase() ===
+                  asset.symbol?.toLocaleLowerCase()
+            )
+          );
+
+          void fetchBalances(fetchBalanceAccounts, {
+            selectedCustomTokens,
+            shouldFetchCustomTokens: selectedCustomTokens.length > 0,
+          });
+        }
       }
 
       setNotification(event, route);

package/src/store/slices/settings.ts

@@ -189,7 +189,7 @@ export const createSettingsSlice: StateCreator<
     });
   },
   setCustomToken: (token) => {
-    void get().fetchCustomTokensBalance({
+    void get().fetchCustomTokensBalances({
       tokens: [token],
       connectedWallets: get().connectedWallets,
     });

package/src/store/slices/wallets.ts

@@ -121,10 +121,17 @@ export interface WalletsSlice {
   ) => void;
   clearConnectedWallet: () => void;
   fetchBalances: (
+    accounts: Wallet[],
+    options?: {
+      shouldFetchCustomTokens?: boolean;
+      selectedCustomTokens?: Asset[];
+    }
+  ) => Promise<void>;
+  fetchMainTokensBalances: (
     accounts: Wallet[],
     options?: { retryOnFailedBalances?: boolean }
   ) => Promise<void>;
-  fetchCustomTokensBalance: (params: {
+  fetchCustomTokensBalances: (params: {
     tokens: Asset[];
     connectedWallets: Wallet[];
   }) => Promise<void>;
@@ -299,7 +306,7 @@ export const createWalletsSlice = keepLastUpdated<AppStoreState, WalletsSlice>(
         });
       }
     },
-    fetchCustomTokensBalance: async (params) => {
+    fetchCustomTokensBalances: async (params) => {
       const { tokens, connectedWallets } = params;
 
       const tokensByBlockchain = tokens.reduce<{ [key: string]: Asset[] }>(
@@ -436,10 +443,6 @@ export const createWalletsSlice = keepLastUpdated<AppStoreState, WalletsSlice>(
       get().addConnectedWallet(accounts, namespace);
 
       void get().fetchBalances(accounts);
-      void get().fetchCustomTokensBalance({
-        tokens: get().customTokens(),
-        connectedWallets: accounts,
-      });
     },
     removeBalancesForWallet: (walletType, options) => {
       const partialCurrentState = {
@@ -581,6 +584,20 @@ export const createWalletsSlice = keepLastUpdated<AppStoreState, WalletsSlice>(
     },
     clearConnectedWallet: () => set({ connectedWallets: [] }),
     fetchBalances: async (accounts, options) => {
+      const { shouldFetchCustomTokens = true, selectedCustomTokens } =
+        options || {};
+      await get().fetchMainTokensBalances(accounts);
+      if (shouldFetchCustomTokens) {
+        void get().fetchCustomTokensBalances({
+          tokens: selectedCustomTokens ?? get().customTokens(),
+          connectedWallets: accounts,
+        });
+      }
+    },
+    fetchMainTokensBalances: async (accounts, options) => {
+      if (accounts.length === 0) {
+        return;
+      }
       // All the `accounts` have same `walletType` so we can pick the first one.
       const walletType = accounts[0].walletType;
 
@@ -602,22 +619,6 @@ export const createWalletsSlice = keepLastUpdated<AppStoreState, WalletsSlice>(
       const walletsDetails = response.wallets;
 
       if (walletsDetails) {
-        const { retryOnFailedBalances = true } = options || {};
-        if (retryOnFailedBalances) {
-          const failedWallets: Wallet[] = walletsDetails
-            .filter((wallet) => wallet.failed)
-            .map((wallet) => ({
-              chain: wallet.blockChain,
-              walletType: walletType,
-              address: wallet.address,
-            }));
-          if (failedWallets.length > 0) {
-            void get().fetchBalances(failedWallets, {
-              retryOnFailedBalances: false,
-            });
-          }
-        }
-
         let nextBalances: BalanceState = get()._balances;
         let nextAggregatedBalances: AggregatedBalanceState =
           get()._aggregatedBalances;
@@ -691,6 +692,27 @@ export const createWalletsSlice = keepLastUpdated<AppStoreState, WalletsSlice>(
         }));
 
         get().setConnectedWalletRetrievedData(accounts, walletsDetails);
+        /*
+         * We handle failed balance fetches after updating the state with successful ones.
+         * This ensures that successful balance updates are immediately available to the UI,
+         * while failed ones are retried in the background without blocking the main flow.
+         * The retry is only attempted once to prevent infinite loops.
+         */
+        const { retryOnFailedBalances = true } = options || {};
+        if (retryOnFailedBalances) {
+          const failedWallets: Wallet[] = walletsDetails
+            .filter((wallet) => wallet.failed)
+            .map((wallet) => ({
+              chain: wallet.blockChain,
+              walletType: walletType,
+              address: wallet.address,
+            }));
+          if (failedWallets.length > 0) {
+            await get().fetchMainTokensBalances(failedWallets, {
+              retryOnFailedBalances: false,
+            });
+          }
+        }
       } else {
         get().setConnectedWalletHasError(accounts);
         console.error(