From e4d6a0419ff1742829d2e63f3ccfa670bd70df08 Mon Sep 17 00:00:00 2001 From: Joaquim Verges Date: Wed, 16 Apr 2025 20:25:13 +1200 Subject: [PATCH] [SDK] Optimize token balance fetching with Insight API --- packages/thirdweb/src/insight/common.ts | 9 +- .../ConnectWallet/screens/Buy/BuyScreen.tsx | 8 +- .../screens/Buy/swap/TokenSelectorScreen.tsx | 126 +++-------- .../Buy/swap/fetchBalancesForWallet.ts | 196 ++++++++++++++++++ 4 files changed, 240 insertions(+), 99 deletions(-) create mode 100644 packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/fetchBalancesForWallet.ts diff --git a/packages/thirdweb/src/insight/common.ts b/packages/thirdweb/src/insight/common.ts index aec7da2e83a..22afc54b7a4 100644 --- a/packages/thirdweb/src/insight/common.ts +++ b/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); +} diff --git a/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/BuyScreen.tsx b/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/BuyScreen.tsx index c7c9445651a..0370652f580 100644 --- a/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/BuyScreen.tsx +++ b/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/BuyScreen.tsx @@ -4,7 +4,10 @@ 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 { + 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 @@ function createSupportedTokens( 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 && diff --git a/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/TokenSelectorScreen.tsx b/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/TokenSelectorScreen.tsx index 30fed0bc9e8..c29e1fe02e8 100644 --- a/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/TokenSelectorScreen.tsx +++ b/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/TokenSelectorScreen.tsx @@ -7,14 +7,8 @@ import { 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 { formatTokenBalance } from "../../formatTokenBalance.js"; 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 @@ export function TokenSelectorScreen(props: { 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(); - - 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(); + for (const entry of entries) { + map.set(entry[0], entry[1]); + } + return map; }, - enabled: !!props.sourceSupportedTokens && !!chainInfo.data, }); if ( diff --git a/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/fetchBalancesForWallet.ts b/packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/swap/fetchBalancesForWallet.ts new file mode 100644 index 00000000000..3233ad2daa5 --- /dev/null +++ b/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(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 { + 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 = { + ...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[] = []; + + 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 = {}; + 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; +}