fix: chain switching with in-app smart accounts (#3005)

Author: joaquim-verges

.changeset/sweet-spoons-move.md

@@ -0,0 +1,7 @@
+---
+"thirdweb": patch
+---
+
+- useSendTransaction now switches chain if needed
+- Fix chain switching with in-app smart accounts
+- Fix tx simulation for pay modal

packages/thirdweb/src/react/core/hooks/contract/useSendTransaction.ts

@@ -5,12 +5,17 @@ import { sendTransaction } from "../../../../transaction/actions/send-transactio
 import type { WaitForReceiptOptions } from "../../../../transaction/actions/wait-for-tx-receipt.js";
 import type { PreparedTransaction } from "../../../../transaction/prepare-transaction.js";
 import { resolvePromisedValue } from "../../../../utils/promise/resolve-promised-value.js";
+import type { Account } from "../../../../wallets/interfaces/wallet.js";
 import {
   type GetWalletBalanceResult,
   getWalletBalance,
 } from "../../../../wallets/utils/getWalletBalance.js";
 import { fetchBuySupportedDestinations } from "../../../web/ui/ConnectWallet/screens/Buy/swap/useSwapSupportedChains.js";
-import { useActiveAccount } from "../wallets/wallet-hooks.js";
+import {
+  useActiveAccount,
+  useActiveWallet,
+  useSwitchActiveWalletChain,
+} from "../wallets/wallet-hooks.js";
 
 type ShowModalData = {
   tx: PreparedTransaction;
@@ -37,10 +42,21 @@ export function useSendTransactionCore(
   showPayModal?: (data: ShowModalData) => void,
   gasless?: GaslessOptions,
 ): UseMutationResult<WaitForReceiptOptions, Error, PreparedTransaction> {
-  const account = useActiveAccount();
+  let _account = useActiveAccount();
+  const wallet = useActiveWallet();
+  const switchChain = useSwitchActiveWalletChain();
 
   return useMutation({
     mutationFn: async (tx) => {
+      // switch chain if needed
+      if (wallet && tx.chain.id !== wallet.getChain()?.id) {
+        await switchChain(tx.chain);
+        // in smart wallet case, account may change after chain switch
+        _account = wallet.getAccount();
+      }
+
+      const account = _account;
+
       if (!account) {
         throw new Error("No active account");
       }
@@ -90,7 +106,7 @@ export function useSendTransactionCore(
                 chain: tx.chain,
                 client: tx.client,
               }),
-              getTotalTxCostForBuy(tx),
+              getTotalTxCostForBuy(tx, account),
             ]);
 
             const walletBalanceWei = walletBalance.value;
@@ -122,9 +138,13 @@ export function useSendTransactionCore(
   });
 }
 
-export async function getTotalTxCostForBuy(tx: PreparedTransaction) {
+export async function getTotalTxCostForBuy(
+  tx: PreparedTransaction,
+  account?: Account,
+) {
   const gasCost = await estimateGasCost({
     transaction: tx,
+    account,
   });
 
   const bufferCost = gasCost.wei / 10n;

packages/thirdweb/src/react/web/ui/ConnectWallet/ConnectButton.tsx

@@ -102,7 +102,7 @@ export function ConnectButton(props: ConnectButtonProps) {
         wallets: wallets,
         locale: props.locale || "en_US",
         connectLocale: localeQuery.data,
-        chain: props.chain,
+        chain: props.chain || props.accountAbstraction?.chain,
         chains: props.chains,
         walletConnect: props.walletConnect,
         accountAbstraction: props.accountAbstraction,

packages/thirdweb/src/react/web/ui/ConnectWallet/Modal/ConnectEmbed.tsx

@@ -389,7 +389,7 @@ export function ConnectEmbed(props: ConnectEmbedProps) {
           wallets: wallets,
           locale: localeId,
           connectLocale: localeQuery.data,
-          chain: props.chain,
+          chain: props.chain || props.accountAbstraction?.chain,
           chains: props.chains,
           walletConnect: props.walletConnect,
           accountAbstraction: props.accountAbstraction,

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

@@ -217,6 +217,7 @@ function BuyScreenContent(props: BuyScreenContentProps) {
     buyForTx,
     hasEditedAmount,
     isMainScreen: screen.type === "main",
+    account,
   });
 
   // check if the screen is expanded or not

packages/thirdweb/src/react/web/ui/ConnectWallet/screens/Buy/main/useBuyTxStates.ts

@@ -1,6 +1,7 @@
 import { useEffect, useState } from "react";
 import { formatNumber } from "../../../../../../../utils/formatNumber.js";
 import { toEther } from "../../../../../../../utils/units.js";
+import type { Account } from "../../../../../../../wallets/interfaces/wallet.js";
 import { getTotalTxCostForBuy } from "../../../../../../core/hooks/contract/useSendTransaction.js";
 import { wait } from "../../../../../../core/utils/wait.js";
 import type { BuyForTx } from "./types.js";
@@ -10,8 +11,10 @@ export function useBuyTxStates(options: {
   buyForTx?: BuyForTx;
   hasEditedAmount: boolean;
   isMainScreen: boolean;
+  account?: Account;
 }) {
-  const { buyForTx, hasEditedAmount, isMainScreen, setTokenAmount } = options;
+  const { buyForTx, hasEditedAmount, isMainScreen, setTokenAmount, account } =
+    options;
   const shouldRefreshTokenAmount = !hasEditedAmount && isMainScreen;
   const stopUpdatingAll = !isMainScreen;
 
@@ -35,7 +38,7 @@ export function useBuyTxStates(options: {
       }
 
       try {
-        const totalCost = await getTotalTxCostForBuy(buyForTx.tx);
+        const totalCost = await getTotalTxCostForBuy(buyForTx.tx, account);
 
         if (!mounted) {
           return;
@@ -64,7 +67,13 @@ export function useBuyTxStates(options: {
     return () => {
       mounted = false;
     };
-  }, [buyForTx, shouldRefreshTokenAmount, setTokenAmount, stopUpdatingAll]);
+  }, [
+    buyForTx,
+    shouldRefreshTokenAmount,
+    setTokenAmount,
+    stopUpdatingAll,
+    account,
+  ]);
 
   return {
     amountNeeded,

packages/thirdweb/src/react/web/ui/TransactionButton/index.tsx

@@ -8,11 +8,7 @@ import {
 import type { PreparedTransaction } from "../../../../transaction/prepare-transaction.js";
 import type { TransactionReceipt } from "../../../../transaction/types.js";
 import { stringify } from "../../../../utils/json.js";
-import {
-  useActiveAccount,
-  useActiveWallet,
-  useSwitchActiveWalletChain,
-} from "../../../core/hooks/wallets/wallet-hooks.js";
+import { useActiveAccount } from "../../../core/hooks/wallets/wallet-hooks.js";
 import {
   type SendTransactionPayModalConfig,
   useSendTransaction,
@@ -131,9 +127,7 @@ export function TransactionButton(props: TransactionButtonProps) {
     ...buttonProps
   } = props;
   const account = useActiveAccount();
-  const wallet = useActiveWallet();
   const [isPending, setIsPending] = useState(false);
-  const switchChain = useSwitchActiveWalletChain();
 
   const sendTransaction = useSendTransaction({
     gasless,
@@ -153,11 +147,6 @@ export function TransactionButton(props: TransactionButtonProps) {
         try {
           setIsPending(true);
           const resolvedTx = await transaction();
-
-          if (wallet && wallet.getChain()?.id !== resolvedTx.chain.id) {
-            await switchChain(resolvedTx.chain);
-          }
-
           const result = await sendTransaction.mutateAsync(resolvedTx);
 
           if (onTransactionSent) {

packages/thirdweb/src/wallets/in-app/core/wallet/index.ts

@@ -25,10 +25,14 @@ export async function connectInAppWallet(
       client: options.client,
       authAccount,
       smartAccountOptions: createOptions.smartAccount,
+      chain: options.chain,
     });
   }
 
-  return [authAccount, options.chain || ethereum] as const;
+  return [
+    authAccount,
+    options.chain || createOptions?.smartAccount?.chain || ethereum,
+  ] as const;
 }
 
 /**
@@ -51,16 +55,21 @@ export async function autoConnectInAppWallet(
       client: options.client,
       authAccount,
       smartAccountOptions: createOptions.smartAccount,
+      chain: options.chain,
     });
   }
 
-  return [authAccount, options.chain || ethereum] as const;
+  return [
+    authAccount,
+    options.chain || createOptions?.smartAccount?.chain || ethereum,
+  ] as const;
 }
 
 async function convertToSmartAccount(options: {
   client: ThirdwebClient;
   authAccount: Account;
   smartAccountOptions: CreateWalletArgs<"smart">[1];
+  chain?: Chain;
 }) {
   const [{ smartWallet }, { connectSmartWallet }] = await Promise.all([
     import("../../../create-wallet.js"),
@@ -73,6 +82,7 @@ async function convertToSmartAccount(options: {
     {
       client: options.client,
       personalAccount: options.authAccount,
+      chain: options.chain,
     },
     options.smartAccountOptions,
   );