fix(sdk): smart wallet account switching

Author: gregfromstl

biome.json

@@ -29,7 +29,8 @@
         "noRestrictedGlobals": {
           "options": { "deniedGlobals": ["Buffer"] },
           "level": "error"
-        }
+        },
+        "noUselessElse": "off"
       }
     }
   },

packages/thirdweb/src/wallets/manager/connection-manager.test.ts

@@ -0,0 +1,68 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { TEST_CLIENT } from "../../../test/src/test-clients.js";
+import { TEST_ACCOUNT_A } from "../../../test/src/test-wallets.js";
+import { sepolia } from "../../chains/chain-definitions/sepolia.js";
+import type { ThirdwebClient } from "../../client/client.js";
+import type { AsyncStorage } from "../../utils/storage/AsyncStorage.js";
+import type { Account, Wallet } from "../interfaces/wallet.js";
+import type { SmartWalletOptions } from "../smart/types.js";
+import { createConnectionManager } from "./index.js";
+
+describe.runIf(process.env.TW_SECRET_KEY)("Connection Manager", () => {
+  let storage: AsyncStorage;
+  let client: ThirdwebClient;
+  let wallet: Wallet;
+  let account: Account;
+  let smartWalletOptions: SmartWalletOptions;
+
+  beforeEach(() => {
+    storage = {
+      getItem: vi.fn(),
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    };
+    client = TEST_CLIENT;
+    account = TEST_ACCOUNT_A;
+    wallet = {
+      id: "wallet-id",
+      getAccount: vi.fn().mockReturnValue(account),
+      subscribe: vi.fn(),
+      disconnect: vi.fn(),
+      switchChain: vi.fn(),
+      getChain: vi.fn().mockReturnValue(sepolia),
+      getConfig: vi.fn(),
+    } as unknown as Wallet;
+    smartWalletOptions = {
+      chain: sepolia,
+    } as SmartWalletOptions;
+  });
+
+  it("connect should handle connection and call onConnect", async () => {
+    const manager = createConnectionManager(storage);
+    const onConnect = vi.fn();
+
+    await manager.connect(wallet, { client, onConnect });
+
+    expect(onConnect).toHaveBeenCalledWith(wallet);
+    expect(storage.setItem).toHaveBeenCalledWith(expect.any(String), wallet.id);
+  });
+
+  it("handleConnection should connect smart wallet", async () => {
+    const manager = createConnectionManager(storage);
+
+    const smartWallet = await manager.handleConnection(wallet, {
+      client,
+      accountAbstraction: smartWalletOptions,
+    });
+
+    expect(manager.activeWalletStore.getValue()).toBe(smartWallet);
+  });
+
+  it("handleConnection should add wallet to connected wallets", async () => {
+    const manager = createConnectionManager(storage);
+
+    await manager.handleConnection(wallet, { client });
+
+    expect(manager.connectedWallets.getValue()).toContain(wallet);
+  });
+});

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

@@ -128,37 +128,54 @@ export function createConnectionManager(storage: AsyncStorage) {
       throw new Error("Can not set a wallet without an account as active");
     }
 
-    const personalWallet = wallet;
-    let activeWallet = personalWallet;
-    const isInAppSmartAccount = hasSmartAccount(wallet);
-    if (options?.accountAbstraction && !isInAppSmartAccount) {
-      activeWallet = smartWallet(options.accountAbstraction);
-      await activeWallet.connect({
-        personalAccount: wallet.getAccount(),
-        client: options.client,
-      });
-    }
+    const activeWallet = await (async () => {
+      if (options?.accountAbstraction && !hasSmartAccount(wallet)) {
+        return await handleSmartWalletConnection(
+          account,
+          options.client,
+          options.accountAbstraction,
+        );
+      } else {
+        return wallet;
+      }
+    })();
 
     // add personal wallet to connected wallets list
-    addConnectedWallet(personalWallet);
+    addConnectedWallet(wallet);
 
-    if (personalWallet.id !== "smart") {
-      await storage.setItem(LAST_ACTIVE_EOA_ID, personalWallet.id);
+    if (wallet.id !== "smart") {
+      await storage.setItem(LAST_ACTIVE_EOA_ID, wallet.id);
     }
 
+    handleSetActiveWallet(activeWallet);
+
+    wallet.subscribe("accountChanged", async () => {
+      connect(wallet, options);
+    });
+
     return activeWallet;
   };
 
+  const handleSmartWalletConnection = async (
+    signer: Account,
+    client: ThirdwebClient,
+    options: SmartWalletOptions,
+  ) => {
+    const wallet = smartWallet(options);
+
+    await wallet.connect({
+      personalAccount: signer,
+      client: client,
+      chain: options.chain,
+    });
+
+    return wallet;
+  };
+
   const connect = async (wallet: Wallet, options?: ConnectManagerOptions) => {
-    // connectedWallet can be either wallet or smartWallet based on
+    // connectedWallet can be either wallet or smartWallet
     const connectedWallet = await handleConnection(wallet, options);
     options?.onConnect?.(connectedWallet);
-    handleSetActiveWallet(connectedWallet);
-    wallet.subscribe("accountChanged", async () => {
-      const newConnectedWallet = await handleConnection(wallet, options);
-      options?.onConnect?.(newConnectedWallet);
-      handleSetActiveWallet(newConnectedWallet);
-    });
     return connectedWallet;
   };