[SDK] Optimize token balance fetching with Insight API

packages/thirdweb/src/insight/common.ts

@@ -4,9 +4,9 @@ import { getChainServices } from "../chains/utils.js";
 export async function assertInsightEnabled(chains: Chain[]) {
   const chainData = await Promise.all(
     chains.map((chain) =>
-      getChainServices(chain).then((services) => ({
+      isInsightEnabled(chain).then((enabled) => ({
         chain,
-        enabled: services.some((c) => c.service === "insight" && c.enabled),
+        enabled,
       })),
     ),
   );
@@ -22,3 +22,8 @@ export async function assertInsightEnabled(chains: Chain[]) {
     );
   }
 }
+
+export async function isInsightEnabled(chain: Chain) {
+  const chainData = await getChainServices(chain);
+  return chainData.some((c) => c.service === "insight" && c.enabled);
+}

packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/BuyScreen.tsx

@@ -4,7 +4,10 @@
 import type { Chain } from "../../../../../../chains/types.js";
 import { getCachedChain } from "../../../../../../chains/utils.js";
 import type { ThirdwebClient } from "../../../../../../client/client.js";
-import { NATIVE_TOKEN_ADDRESS } from "../../../../../../constants/addresses.js";
+import {
+  NATIVE_TOKEN_ADDRESS,
+  ZERO_ADDRESS,
+} from "../../../../../../constants/addresses.js";
 import type { BuyWithCryptoStatus } from "../../../../../../pay/buyWithCrypto/getStatus.js";
 import type { BuyWithFiatStatus } from "../../../../../../pay/buyWithFiat/getStatus.js";
 import { formatNumber } from "../../../../../../utils/formatNumber.js";
@@ -982,6 +985,9 @@
 
   for (const x of data) {
     tokens[x.chain.id] = x.tokens.filter((t) => {
+      if (t.address === ZERO_ADDRESS) {
+        return false;
+      }
       // for source tokens, data is not provided, so we include all of them
       if (
         t.buyWithCryptoEnabled === undefined &&

packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/TokenSelectorScreen.tsx

@@ -7,14 +7,8 @@
 import { useQuery } from "@tanstack/react-query";
 import { trackPayEvent } from "../../../../../../../analytics/track/pay.js";
 import type { Chain } from "../../../../../../../chains/types.js";
-import { getCachedChain } from "../../../../../../../chains/utils.js";
 import type { ThirdwebClient } from "../../../../../../../client/client.js";
-import { NATIVE_TOKEN_ADDRESS } from "../../../../../../../constants/addresses.js";
 import type { Wallet } from "../../../../../../../wallets/interfaces/wallet.js";
-import {
-  type GetWalletBalanceResult,
-  getWalletBalance,
-} from "../../../../../../../wallets/utils/getWalletBalance.js";
 import type { WalletId } from "../../../../../../../wallets/wallet-types.js";
 import { useCustomTheme } from "../../../../../../core/design-system/CustomThemeProvider.js";
 import {
@@ -47,12 +41,10 @@
 import { type ERC20OrNativeToken, isNativeToken } from "../../nativeToken.js";
 import { FiatValue } from "./FiatValue.js";
 import { WalletRow } from "./WalletRow.js";
-
-type TokenBalance = {
-  balance: GetWalletBalanceResult;
-  chain: Chain;
-  token: TokenInfo;
-};
+import {
+  type TokenBalance,
+  fetchBalancesForWallet,
+} from "./fetchBalancesForWallet.js";
 
 type WalletKey = {
   id: WalletId;
@@ -90,93 +82,35 @@
       activeAccount?.address,
       connectedWallets.map((w) => w.getAccount()?.address),
     ],
+    enabled: !!props.sourceSupportedTokens && !!chainInfo.data,
     queryFn: async () => {
-      // in parallel, get the balances of all the wallets on each of the sourceSupportedTokens
-      const walletBalanceMap = new Map<WalletKey, TokenBalance[]>();
-
-      const balancePromises = connectedWallets.flatMap((wallet) => {
-        const account = wallet.getAccount();
-        if (!account) return [];
-        const walletKey: WalletKey = {
-          id: wallet.id,
-          address: account.address,
-        };
-        walletBalanceMap.set(walletKey, []);
-
-        // inject the destination token too since it can be used as well to pay/transfer
-        const toToken = isNativeToken(props.toToken)
-          ? {
-              address: NATIVE_TOKEN_ADDRESS,
-              name: chainInfo.data?.nativeCurrency.name || "",
-              symbol: chainInfo.data?.nativeCurrency.symbol || "",
-              icon: chainInfo.data?.icon?.url,
-            }
-          : props.toToken;
-
-        const tokens = {
-          ...props.sourceSupportedTokens,
-          [props.toChain.id]: [
-            toToken,
-            ...(props.sourceSupportedTokens?.[props.toChain.id] || []),
-          ],
-        };
-
-        return Object.entries(tokens).flatMap(([chainId, tokens]) => {
-          return tokens.map(async (token) => {
-            try {
-              const chain = getCachedChain(Number(chainId));
-              const balance = await getWalletBalance({
-                address: account.address,
-                chain,
-                tokenAddress: isNativeToken(token) ? undefined : token.address,
-                client: props.client,
-              });
-
-              // show the token if:
-              // - its not the destination token and balance is greater than 0
-              // - its the destination token and balance is greater than the token amount AND we the account is not the default account in fund_wallet mode
-              const shouldInclude =
-                token.address === toToken.address &&
-                chain.id === props.toChain.id
-                  ? props.mode === "fund_wallet" &&
-                    account.address === activeAccount?.address
-                    ? false
-                    : Number(balance.displayValue) > Number(props.tokenAmount)
-                  : balance.value > 0n;
-
-              if (shouldInclude) {
-                const existingBalances = walletBalanceMap.get(walletKey) || [];
-                existingBalances.push({ balance, chain, token });
-                existingBalances.sort((a, b) => {
-                  if (
-                    a.chain.id === props.toChain.id &&
-                    a.token.address === toToken.address
-                  )
-                    return -1;
-                  if (
-                    b.chain.id === props.toChain.id &&
-                    b.token.address === toToken.address
-                  )
-                    return 1;
-                  if (a.chain.id === props.toChain.id) return -1;
-                  if (b.chain.id === props.toChain.id) return 1;
-                  return a.chain.id > b.chain.id ? 1 : -1;
-                });
-              }
-            } catch (error) {
-              console.error(
-                `Failed to fetch balance for wallet ${wallet.id} on chain ${chainId} for token ${token.symbol}:`,
-                error,
-              );
-            }
+      const entries = await Promise.all(
+        connectedWallets.map(async (wallet) => {
+          const balances = await fetchBalancesForWallet({
+            wallet,
+            accountAddress: activeAccount?.address,
+            sourceSupportedTokens: props.sourceSupportedTokens || [],
+            toChain: props.toChain,
+            toToken: props.toToken,
+            tokenAmount: props.tokenAmount,
+            mode: props.mode,
+            client: props.client,
           });
-        });
-      });
-
-      await Promise.all(balancePromises);
-      return walletBalanceMap;
+          return [
+            {
+              id: wallet.id,
+              address: wallet.getAccount()?.address || "",
+            } as WalletKey,
+            balances,
+          ] as const;
+        }),
+      );
+      const map = new Map<WalletKey, TokenBalance[]>();
+      for (const entry of entries) {
+        map.set(entry[0], entry[1]);
+      }
+      return map;
     },
-    enabled: !!props.sourceSupportedTokens && !!chainInfo.data,
   });
 
   if (

packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/fetchBalancesForWallet.ts

@@ -0,0 +1,196 @@
+import type { Chain } from "../../../../../../../chains/types.js";
+import { getCachedChain } from "../../../../../../../chains/utils.js";
+import type { ThirdwebClient } from "../../../../../../../client/client.js";
+import { NATIVE_TOKEN_ADDRESS } from "../../../../../../../constants/addresses.js";
+import { isInsightEnabled } from "../../../../../../../insight/common.js";
+import { getOwnedTokens } from "../../../../../../../insight/get-tokens.js";
+import type { Wallet } from "../../../../../../../wallets/interfaces/wallet.js";
+import {
+  type GetWalletBalanceResult,
+  getWalletBalance,
+} from "../../../../../../../wallets/utils/getWalletBalance.js";
+import type { PayUIOptions } from "../../../../../../core/hooks/connection/ConnectButtonProps.js";
+import type {
+  SupportedTokens,
+  TokenInfo,
+} from "../../../../../../core/utils/defaultTokens.js";
+import { type ERC20OrNativeToken, isNativeToken } from "../../nativeToken.js";
+
+const CHUNK_SIZE = 5;
+
+function chunkChains<T>(chains: T[]): T[][] {
+  const chunks: T[][] = [];
+  for (let i = 0; i < chains.length; i += CHUNK_SIZE) {
+    chunks.push(chains.slice(i, i + CHUNK_SIZE));
+  }
+  return chunks;
+}
+
+type FetchBalancesParams = {
+  wallet: Wallet;
+  accountAddress: string | undefined;
+  sourceSupportedTokens: SupportedTokens;
+  toChain: Chain;
+  toToken: ERC20OrNativeToken;
+  tokenAmount: string;
+  mode: PayUIOptions["mode"];
+  client: ThirdwebClient;
+};
+
+export type TokenBalance = {
+  balance: GetWalletBalanceResult;
+  chain: Chain;
+  token: TokenInfo;
+};
+
+export async function fetchBalancesForWallet({
+  wallet,
+  accountAddress,
+  sourceSupportedTokens,
+  toChain,
+  toToken,
+  tokenAmount,
+  mode,
+  client,
+}: FetchBalancesParams): Promise<TokenBalance[]> {
+  const account = wallet.getAccount();
+  if (!account) {
+    return [];
+  }
+
+  const balances: TokenBalance[] = [];
+
+  // 1. Resolve all unique chains in the supported token map
+  const uniqueChains = Object.keys(sourceSupportedTokens).map((id) =>
+    getCachedChain(Number(id)),
+  );
+
+  // 2. Check insight availability once per chain
+  const insightSupport = await Promise.all(
+    uniqueChains.map(async (c) => ({
+      chain: c,
+      enabled: await isInsightEnabled(c),
+    })),
+  );
+  const insightEnabledChains = insightSupport
+    .filter((c) => c.enabled)
+    .map((c) => c.chain);
+
+  // 3. ERC-20 balances for insight-enabled chains (batched 5 chains / call)
+  const insightChunks = chunkChains(insightEnabledChains);
+  await Promise.all(
+    insightChunks.map(async (chunk) => {
+      const owned = await getOwnedTokens({
+        ownerAddress: account.address,
+        chains: chunk,
+        client,
+      });
+
+      for (const b of owned) {
+        const matching = sourceSupportedTokens[b.chainId]?.find(
+          (t) => t.address.toLowerCase() === b.tokenAddress.toLowerCase(),
+        );
+        if (matching) {
+          balances.push({
+            balance: b,
+            chain: getCachedChain(b.chainId),
+            token: matching,
+          });
+        }
+      }
+    }),
+  );
+
+  // 4. Build a token map that also includes the destination token so it can be used to pay
+  const destinationToken = isNativeToken(toToken)
+    ? {
+        address: NATIVE_TOKEN_ADDRESS,
+        name: toChain.nativeCurrency?.name || "",
+        symbol: toChain.nativeCurrency?.symbol || "",
+        icon: toChain.icon?.url,
+      }
+    : toToken;
+
+  const tokenMap: Record<number, TokenInfo[]> = {
+    ...sourceSupportedTokens,
+    [toChain.id]: [
+      destinationToken,
+      ...(sourceSupportedTokens[toChain.id] || []),
+    ],
+  };
+
+  // 5. Fallback RPC balances (native currency & ERC-20 that we couldn't fetch from insight)
+  const rpcCalls: Promise<void>[] = [];
+
+  for (const [chainIdStr, tokens] of Object.entries(tokenMap)) {
+    const chainId = Number(chainIdStr);
+    const chain = getCachedChain(chainId);
+
+    for (const token of tokens) {
+      const isNative = isNativeToken(token);
+      const isAlreadyFetched = balances.some(
+        (b) =>
+          b.chain.id === chainId &&
+          b.token.address.toLowerCase() === token.address.toLowerCase(),
+      );
+      if (!isNative && !isAlreadyFetched) {
+        // ERC20 on insight-enabled chain already handled by insight call
+        continue;
+      }
+      rpcCalls.push(
+        (async () => {
+          try {
+            const balance = await getWalletBalance({
+              address: account.address,
+              chain,
+              tokenAddress: isNative ? undefined : token.address,
+              client,
+            });
+
+            const include =
+              token.address === destinationToken.address &&
+              chain.id === toChain.id
+                ? mode === "fund_wallet" && account.address === accountAddress
+                  ? false
+                  : Number(balance.displayValue) > Number(tokenAmount)
+                : balance.value > 0n;
+
+            if (include) {
+              balances.push({ balance, chain, token });
+            }
+          } catch (err) {
+            console.warn(
+              `Failed to fetch balance for ${token.symbol} on chain ${chainId}`,
+              err,
+            );
+          }
+        })(),
+      );
+    }
+  }
+
+  await Promise.all(rpcCalls);
+
+  // Remove duplicates (same chainId + token address)
+  {
+    const uniq: Record<string, TokenBalance> = {};
+    for (const b of balances) {
+      const k = `${b.chain.id}-${b.token.address.toLowerCase()}`;
+      if (!uniq[k]) {
+        uniq[k] = b;
+      }
+    }
+    balances.splice(0, balances.length, ...Object.values(uniq));
+  }
+  // 6. Sort so that the destination token always appears first, then tokens on the destination chain, then by chain id
+  balances.sort((a, b) => {
+    const destAddress = destinationToken.address;
+    if (a.chain.id === toChain.id && a.token.address === destAddress) return -1;
+    if (b.chain.id === toChain.id && b.token.address === destAddress) return 1;
+    if (a.chain.id === toChain.id) return -1;
+    if (b.chain.id === toChain.id) return 1;
+    return a.chain.id - b.chain.id;
+  });
+
+  return balances;
+}