Preload wallet balances on pay embed

Author: joaquim-verges

.changeset/sour-flies-post.md

@@ -0,0 +1,5 @@
+---
+"thirdweb": patch
+---
+
+Preload wallet balances on pay embed

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

@@ -63,6 +63,7 @@ import { SwapFlow } from "./swap/SwapFlow.js";
 import { SwapScreenContent } from "./swap/SwapScreenContent.js";
 import { TokenSelectorScreen } from "./swap/TokenSelectorScreen.js";
 import { TransferFlow } from "./swap/TransferFlow.js";
+import { useWalletsAndBalances } from "./swap/fetchBalancesForWallet.js";
 import {
   type SupportedChainAndTokens,
   useBuySupportedDestinations,
@@ -218,6 +219,16 @@ function BuyScreenContent(props: BuyScreenContentProps) {
     );
   }, [props.supportedTokens, supportedSourcesQuery.data, payOptions]);
 
+  // preload wallets and balances
+  useWalletsAndBalances({
+    client: props.client,
+    sourceSupportedTokens: sourceSupportedTokens || [],
+    toChain: toChain,
+    toToken: toToken,
+    tokenAmount: tokenAmount,
+    mode: payOptions.mode,
+  });
+
   const { fromChain, setFromChain, fromToken, setFromToken } =
     useFromTokenSelectionStates({
       payOptions,

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

@@ -4,7 +4,6 @@ import {
   ChevronRightIcon,
   Cross2Icon,
 } from "@radix-ui/react-icons";
-import { useQuery } from "@tanstack/react-query";
 import { trackPayEvent } from "../../../../../../../analytics/track/pay.js";
 import type { Chain } from "../../../../../../../chains/types.js";
 import type { ThirdwebClient } from "../../../../../../../client/client.js";
@@ -43,14 +42,9 @@ import { FiatValue } from "./FiatValue.js";
 import { WalletRow } from "./WalletRow.js";
 import {
   type TokenBalance,
-  fetchBalancesForWallet,
+  useWalletsAndBalances,
 } from "./fetchBalancesForWallet.js";
 
-type WalletKey = {
-  id: WalletId;
-  address: string;
-};
-
 export function TokenSelectorScreen(props: {
   client: ThirdwebClient;
   sourceTokens: SupportedTokens | undefined;
@@ -71,46 +65,13 @@ export function TokenSelectorScreen(props: {
   const chainInfo = useChainMetadata(props.toChain);
   const theme = useCustomTheme();
 
-  const walletsAndBalances = useQuery({
-    queryKey: [
-      "wallets-and-balances",
-      props.sourceSupportedTokens,
-      props.toChain.id,
-      props.toToken,
-      props.tokenAmount,
-      props.mode,
-      activeAccount?.address,
-      connectedWallets.map((w) => w.getAccount()?.address),
-    ],
-    enabled: !!props.sourceSupportedTokens && !!chainInfo.data,
-    queryFn: async () => {
-      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,
-          });
-          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;
-    },
+  const walletsAndBalances = useWalletsAndBalances({
+    client: props.client,
+    sourceSupportedTokens: props.sourceSupportedTokens || [],
+    toChain: props.toChain,
+    toToken: props.toToken,
+    tokenAmount: props.tokenAmount,
+    mode: props.mode,
   });
 
   if (

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

@@ -1,3 +1,4 @@
+import { useQuery } from "@tanstack/react-query";
 import type { Chain } from "../../../../../../../chains/types.js";
 import { getCachedChain } from "../../../../../../../chains/utils.js";
 import type { ThirdwebClient } from "../../../../../../../client/client.js";
@@ -9,7 +10,11 @@ import {
   type GetWalletBalanceResult,
   getWalletBalance,
 } from "../../../../../../../wallets/utils/getWalletBalance.js";
+import type { WalletId } from "../../../../../../../wallets/wallet-types.js";
 import type { PayUIOptions } from "../../../../../../core/hooks/connection/ConnectButtonProps.js";
+import { useChainMetadata } from "../../../../../../core/hooks/others/useChainQuery.js";
+import { useActiveAccount } from "../../../../../../core/hooks/wallets/useActiveAccount.js";
+import { useConnectedWallets } from "../../../../../../core/hooks/wallets/useConnectedWallets.js";
 import type {
   SupportedTokens,
   TokenInfo,
@@ -43,7 +48,68 @@ export type TokenBalance = {
   token: TokenInfo;
 };
 
-export async function fetchBalancesForWallet({
+type WalletKey = {
+  id: WalletId;
+  address: string;
+};
+
+export function useWalletsAndBalances(props: {
+  sourceSupportedTokens: SupportedTokens;
+  toChain: Chain;
+  toToken: ERC20OrNativeToken;
+  tokenAmount: string;
+  mode: PayUIOptions["mode"];
+  client: ThirdwebClient;
+}) {
+  const activeAccount = useActiveAccount();
+  const connectedWallets = useConnectedWallets();
+  const chainInfo = useChainMetadata(props.toChain);
+
+  return useQuery({
+    queryKey: [
+      "wallets-and-balances",
+      props.sourceSupportedTokens,
+      props.toChain.id,
+      props.toToken,
+      props.tokenAmount,
+      props.mode,
+      activeAccount?.address,
+      connectedWallets.map((w) => w.getAccount()?.address),
+    ],
+    enabled:
+      !!props.sourceSupportedTokens && !!chainInfo.data && !!activeAccount,
+    queryFn: async () => {
+      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,
+          });
+          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;
+    },
+  });
+}
+
+async function fetchBalancesForWallet({
   wallet,
   accountAddress,
   sourceSupportedTokens,