Fix custom chain not used in wallet connection (#2768)

Co-authored-by: Jonas Daniels <jonas.daniels@outlook.com>

Author: MananTank

.changeset/tiny-cougars-crash.md

@@ -0,0 +1,5 @@
+---
+"thirdweb": patch
+---
+
+Fix custom chain not being used

packages/thirdweb/src/react/core/hooks/connection/useAutoConnect.tsx

@@ -10,6 +10,7 @@ import type {
   // WalletWithPersonalAccount,
 } from "../../../../wallets/interfaces/wallet.js";
 import {
+  getLastConnectedChain,
   getStoredActiveWalletId,
   getStoredConnectedWalletIds,
 } from "../../../../wallets/manager/index.js";
@@ -178,10 +179,13 @@ export function AutoConnect(props: AutoConnectProps) {
         return;
       }
 
+      const lastConnectedChain = await getLastConnectedChain(asyncLocalStorage);
+
       async function handleWalletConnection(wallet: Wallet) {
         setConnectionStatus("connecting");
         return wallet.autoConnect({
           client: props.client,
+          chain: lastConnectedChain ?? undefined,
         });
       }
 

packages/thirdweb/src/wallets/coinbase/coinbaseSDKWallet.ts

@@ -280,7 +280,10 @@ export async function connectCoinbaseWalletSDK(
   })) as string | number;
 
   const chainId = normalizeChainId(connectedChainId);
-  let chain = defineChain(chainId);
+  let chain =
+    options.chain && options.chain.id === chainId
+      ? options.chain
+      : defineChain(chainId);
   // Switch to chain if provided
   if (
     connectedChainId &&
@@ -318,7 +321,10 @@ export async function autoConnectCoinbaseWalletSDK(
     method: "eth_chainId",
   })) as string | number;
   const chainId = normalizeChainId(connectedChainId);
-  const chain = defineChain(chainId);
+  const chain =
+    options.chain && options.chain.id === chainId
+      ? options.chain
+      : defineChain(chainId);
 
   return onConnect(address, chain, provider, emitter);
 }

packages/thirdweb/src/wallets/create-wallet.ts

@@ -8,6 +8,7 @@ import type { Account, Wallet } from "./interfaces/wallet.js";
 import type {
   CreateWalletArgs,
   InjectedConnectOptions,
+  WalletAutoConnectionOption,
   WalletId,
 } from "./wallet-types.js";
 
@@ -104,7 +105,11 @@ export function createWallet<const ID extends WalletId>(
         getConfig: () => args[1],
         getChain: () => chain,
         getAccount: () => account,
-        autoConnect: async (options) => {
+        autoConnect: async (
+          options: WalletAutoConnectionOption<
+            WCSupportedWalletIds | InjectedSupportedWalletIds
+          >,
+        ) => {
           // injected wallet priority for autoConnect
           if (id !== "walletConnect" && injectedProvider(id)) {
             const { autoConnectInjectedWallet } = await import(
@@ -119,6 +124,7 @@ export function createWallet<const ID extends WalletId>(
             ] = await autoConnectInjectedWallet(
               id as InjectedSupportedWalletIds,
               emitter,
+              options.chain,
             );
             // set the states
             account = connectedAccount;

packages/thirdweb/src/wallets/injected/index.ts

@@ -63,7 +63,10 @@ export async function connectInjectedWallet(
     .request({ method: "eth_chainId" })
     .then(normalizeChainId);
 
-  let connectedChain = defineChain(chainId);
+  let connectedChain =
+    options.chain && options.chain.id === chainId
+      ? options.chain
+      : defineChain(chainId);
 
   // if we want a specific chainId and it is not the same as the provider chainId, trigger switchChain
   if (options.chain && options.chain.id !== chainId) {
@@ -80,6 +83,7 @@ export async function connectInjectedWallet(
 export async function autoConnectInjectedWallet(
   id: InjectedSupportedWalletIds,
   emitter: WalletEmitter<InjectedSupportedWalletIds>,
+  chain?: Chain,
 ): Promise<ReturnType<typeof onConnect>> {
   const provider = getInjectedProvider(id);
 
@@ -101,7 +105,8 @@ export async function autoConnectInjectedWallet(
     .request({ method: "eth_chainId" })
     .then(normalizeChainId);
 
-  const connectedChain = defineChain(chainId);
+  const connectedChain =
+    chain && chain.id === chainId ? chain : defineChain(chainId);
 
   return onConnect(provider, address, connectedChain, emitter);
 }

packages/thirdweb/src/wallets/manager/index.ts

@@ -11,6 +11,7 @@ export type ConnectionStatus = "connected" | "disconnected" | "connecting";
 
 const CONNECTED_WALLET_IDS = "thirdweb:connected-wallet-ids";
 const ACTIVE_WALLET_ID = "thirdweb:active-wallet-id";
+const LAST_ACTIVE_CHAIN = "thirdweb:active-chain";
 
 export type ConnectionManager = ReturnType<typeof createConnectionManager>;
 
@@ -126,6 +127,19 @@ export function createConnectionManager(storage: AsyncStorage) {
 
   // side effects
 
+  effect(
+    () => {
+      const _chain = activeWalletChainStore.getValue();
+      if (_chain) {
+        storage.setItem(LAST_ACTIVE_CHAIN, JSON.stringify(_chain));
+      } else {
+        storage.removeItem(LAST_ACTIVE_CHAIN);
+      }
+    },
+    [activeWalletChainStore],
+    false,
+  );
+
   // save last connected wallet ids to storage
   effect(
     () => {
@@ -218,3 +232,21 @@ export async function getStoredActiveWalletId(
     return null;
   }
 }
+
+/**
+ * @internal
+ */
+export async function getLastConnectedChain(
+  storage: AsyncStorage,
+): Promise<Chain | null> {
+  try {
+    const value = await storage.getItem(LAST_ACTIVE_CHAIN);
+    if (value) {
+      return JSON.parse(value) as Chain;
+    }
+
+    return null;
+  } catch {
+    return null;
+  }
+}

packages/thirdweb/src/wallets/wallet-connect/index.ts

@@ -102,7 +102,12 @@ export async function connectWC(
     throw new Error("No accounts found on provider.");
   }
 
-  const chain = defineChain(normalizeChainId(provider.chainId));
+  const providerChainId = normalizeChainId(provider.chainId);
+
+  const chain =
+    options.chain && options.chain.id === providerChainId
+      ? options.chain
+      : defineChain(providerChainId);
 
   if (options) {
     const savedParams: SavedConnectParams = {
@@ -160,7 +165,12 @@ export async function autoConnectWC(
     throw new Error("No accounts found on provider.");
   }
 
-  const chain = defineChain(normalizeChainId(provider.chainId));
+  const providerChainId = normalizeChainId(provider.chainId);
+
+  const chain =
+    options.chain && options.chain.id === providerChainId
+      ? options.chain
+      : defineChain(providerChainId);
 
   return onConnect(address, chain, provider, emitter);
 }

packages/thirdweb/src/wallets/wallet-connect/types.ts

@@ -179,6 +179,8 @@ export type WCAutoConnectOptions = {
   client: ThirdwebClient;
 
   savedConnectParams?: SavedConnectParams;
+
+  chain?: Chain;
 };
 
 type SavedConnectParams = {