[SDK] Fix: Switch underlying admin wallet with ecosystem smart account (#6119)

<!--

## title your PR with this format: "[SDK/Dashboard/Portal] Feature/Fix: Concise title for the changes"

If you did not copy the branch name from Linear, paste the issue tag here (format is TEAM-0000):

## Notes for the reviewer

Anything important to call out? Be sure to also clarify these in your comments.

## How to test

Unit tests, playground, etc.

-->

<!-- start pr-codex -->

---

## PR-Codex overview
This PR focuses on enhancing the functionality of smart wallets in the `thirdweb` ecosystem, ensuring they correctly trigger chain switches and improving related wallet checks and configurations.

### Detailed summary
- Fixed smart wallets to trigger chain switch on admin wallets.
- Updated `smartWalletGetCapabilities` to accept a generic `Wallet` type.
- Implemented `isSmartWallet` function to check for smart wallet status.
- Modified `getSmartWallet` to throw an error for non-smart wallets.
- Refactored wallet checks in `createConnectionManager`.
- Added tests for `isSmartWallet` and `getSmartWallet`.

> ✨ Ask PR-Codex anything about this PR by commenting with `/codex {your question}`

<!-- end pr-codex -->

Author: gregfromstl

.changeset/metal-icons-end.md

@@ -0,0 +1,5 @@
+---
+"thirdweb": patch
+---
+
+Fix: Ecosystem smart wallets now properly trigger switch chain on their admin wallets

packages/thirdweb/src/wallets/eip5792/get-capabilities.ts

@@ -4,7 +4,6 @@ import { isCoinbaseSDKWallet } from "../coinbase/coinbase-web.js";
 import { isInAppWallet } from "../in-app/core/wallet/index.js";
 import { getInjectedProvider } from "../injected/index.js";
 import type { Wallet } from "../interfaces/wallet.js";
-import { isSmartWallet } from "../smart/index.js";
 import { isWalletConnect } from "../wallet-connect/controller.js";
 import type { WalletId } from "../wallet-types.js";
 import type { WalletCapabilities, WalletCapabilitiesRecord } from "./types.js";
@@ -47,7 +46,7 @@ export async function getCapabilities<const ID extends WalletId = WalletId>({
     };
   }
 
-  if (isSmartWallet(wallet)) {
+  if (wallet.id === "smart") {
     const { smartWalletGetCapabilities } = await import(
       "../smart/lib/smart-wallet-capabilities.js"
     );

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

@@ -1,10 +1,13 @@
 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 { baseSepolia } from "../../chains/chain-definitions/base-sepolia.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 { inAppWallet } from "../in-app/web/in-app.js";
 import type { Account, Wallet } from "../interfaces/wallet.js";
+import { smartWallet } from "../smart/smart-wallet.js";
 import type { SmartWalletOptions } from "../smart/types.js";
 import {
   createConnectionManager,
@@ -157,6 +160,46 @@ describe.runIf(process.env.TW_SECRET_KEY)("Connection Manager", () => {
     expect(wallet.switchChain).toHaveBeenCalledWith(newChain);
   });
 
+  it("should switch admin wallet for smart wallet if available", async () => {
+    const manager = createConnectionManager(storage);
+    const adminAccount = TEST_ACCOUNT_A;
+    const adminWallet = inAppWallet();
+    adminWallet.getAccount = () => adminAccount;
+
+    const _wallet = smartWallet({
+      chain: baseSepolia,
+      sponsorGas: true,
+    });
+    await _wallet.connect({
+      client,
+      personalAccount: adminAccount,
+    });
+
+    await manager.handleConnection(adminWallet, { client });
+    await manager.handleConnection(_wallet, { client });
+
+    const newChain = {
+      id: 2,
+      name: "New Chain",
+      rpc: "https://rpc.example.com",
+    };
+
+    // Mock storage and wallet setup
+    storage.getItem = vi.fn().mockResolvedValue("inApp");
+    adminWallet.id = "inApp";
+    _wallet.switchChain = vi.fn();
+    adminWallet.switchChain = vi.fn();
+
+    // Add wallets to connected wallets store
+    manager.addConnectedWallet(adminWallet);
+    manager.addConnectedWallet(_wallet);
+
+    await manager.switchActiveWalletChain(newChain);
+
+    expect(_wallet.switchChain).toHaveBeenCalledWith(newChain);
+    expect(adminWallet.switchChain).toHaveBeenCalledWith(newChain);
+  });
+
   it("should define chains", async () => {
     const manager = createConnectionManager(storage);
     await manager.handleConnection(wallet, { client });

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

@@ -9,6 +9,7 @@ import { stringify } from "../../utils/json.js";
 import type { AsyncStorage } from "../../utils/storage/AsyncStorage.js";
 import { deleteConnectParamsFromStorage } from "../../utils/storage/walletStorage.js";
 import type { Account, Wallet } from "../interfaces/wallet.js";
+import { isSmartWallet } from "../smart/index.js";
 import { smartWallet } from "../smart/smart-wallet.js";
 import type { SmartWalletOptions } from "../smart/types.js";
 import type { WalletId } from "../wallet-types.js";
@@ -257,7 +258,7 @@ export function createConnectionManager(storage: AsyncStorage) {
       throw new Error("Wallet does not support switching chains");
     }
 
-    if (wallet.id === "smart") {
+    if (isSmartWallet(wallet)) {
       // also switch personal wallet
       const personalWalletId = await getStoredActiveWalletId(storage);
       if (personalWalletId) {
@@ -266,10 +267,14 @@ export function createConnectionManager(storage: AsyncStorage) {
           .find((w) => w.id === personalWalletId);
         if (personalWallet) {
           await personalWallet.switchChain(chain);
+          await wallet.switchChain(chain);
+          // reset the active wallet as switch chain recreates a new smart account
+          handleSetActiveWallet(wallet);
+          return;
         }
       }
+      // If we couldn't find the personal wallet, just switch the smart wallet
       await wallet.switchChain(chain);
-      // reset the active wallet as switch chain recreates a new smart account
       handleSetActiveWallet(wallet);
     } else {
       await wallet.switchChain(chain);

packages/thirdweb/src/wallets/smart/get-smart-wallet-config.test.ts

@@ -0,0 +1,67 @@
+import { describe, expect, it } from "vitest";
+import { optimism } from "../../chains/chain-definitions/optimism.js";
+import type { Wallet } from "../interfaces/wallet.js";
+import { getSmartWallet } from "./get-smart-wallet-config.js";
+import type { SmartWalletOptions } from "./types.js";
+
+describe("getSmartWallet", () => {
+  const mockSmartWalletConfig: SmartWalletOptions = {
+    chain: optimism,
+    sponsorGas: false,
+  };
+
+  it("should return config for smart wallet ID", () => {
+    const wallet = {
+      id: "smart",
+      getConfig: () => mockSmartWalletConfig,
+    } as Wallet<"smart">;
+
+    expect(getSmartWallet(wallet)).toBe(mockSmartWalletConfig);
+  });
+
+  it("should return smartAccount config for wallet with smartAccount", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => ({
+        smartAccount: mockSmartWalletConfig,
+      }),
+    } as Wallet;
+
+    expect(getSmartWallet(wallet)).toBe(mockSmartWalletConfig);
+  });
+
+  it("should throw error for non-smart wallet", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => ({}),
+    } as Wallet;
+
+    expect(() => getSmartWallet(wallet)).toThrow(
+      "Wallet is not a smart wallet",
+    );
+  });
+
+  it("should throw error when getConfig returns null", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => null,
+      // biome-ignore lint/suspicious/noExplicitAny: Testing invalid config
+    } as any as Wallet;
+
+    expect(() => getSmartWallet(wallet)).toThrow(
+      "Wallet is not a smart wallet",
+    );
+  });
+
+  it("should throw error when smartAccount is null", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => ({ smartAccount: null }),
+      // biome-ignore lint/suspicious/noExplicitAny: Testing invalid config
+    } as any as Wallet;
+
+    expect(() => getSmartWallet(wallet)).toThrow(
+      "Wallet is not a smart wallet",
+    );
+  });
+});

packages/thirdweb/src/wallets/smart/get-smart-wallet-config.ts

@@ -0,0 +1,24 @@
+import type { Wallet } from "../interfaces/wallet.js";
+import type { SmartWalletOptions } from "./types.js";
+
+/**
+ * Gets the smart wallet configuration for a given wallet.
+ *
+ * @param {Wallet} wallet - The wallet to check.
+ * @returns {SmartWalletOptions} The smart wallet configuration.
+ *
+ * @throws {Error} If the wallet is not a smart wallet.
+ * @internal
+ */
+export function getSmartWallet(wallet: Wallet): SmartWalletOptions {
+  if (wallet.id === "smart") {
+    return (wallet as Wallet<"smart">).getConfig();
+  }
+
+  const config = wallet.getConfig();
+  if (!!config && "smartAccount" in config && !!config?.smartAccount) {
+    return config.smartAccount;
+  }
+
+  throw new Error("Wallet is not a smart wallet");
+}

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

@@ -20,12 +20,7 @@ import type { Hex } from "../../utils/encoding/hex.js";
 import { resolvePromisedValue } from "../../utils/promise/resolve-promised-value.js";
 import { parseTypedData } from "../../utils/signatures/helpers/parse-typed-data.js";
 import { type SignableMessage, maxUint96 } from "../../utils/types.js";
-import type {
-  Account,
-  SendTransactionOption,
-  Wallet,
-} from "../interfaces/wallet.js";
-import type { WalletId } from "../wallet-types.js";
+import type { Account, SendTransactionOption } from "../interfaces/wallet.js";
 import {
   broadcastZkTransaction,
   bundleUserOp,
@@ -58,17 +53,7 @@ import type {
   UserOperationV06,
   UserOperationV07,
 } from "./types.js";
-/**
- * Checks if the provided wallet is a smart wallet.
- *
- * @param wallet - The wallet to check.
- * @returns True if the wallet is a smart wallet, false otherwise.
- */
-export function isSmartWallet(
-  wallet: Wallet<WalletId>,
-): wallet is Wallet<"smart"> {
-  return wallet.id === "smart";
-}
+export { isSmartWallet } from "./is-smart-wallet.js";
 
 /**
  * For in-app wallets, the smart wallet creation is implicit so we track these to be able to retrieve the personal account for a smart account on the wallet API.

packages/thirdweb/src/wallets/smart/is-smart-wallet.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, it } from "vitest";
+import type { Wallet } from "../interfaces/wallet.js";
+import { isSmartWallet } from "./is-smart-wallet.js";
+
+describe("isSmartWallet", () => {
+  it("should return true for smart wallet ID", () => {
+    const wallet = {
+      id: "smart",
+    } as Wallet;
+    expect(isSmartWallet(wallet)).toBe(true);
+  });
+
+  it("should return true for wallet with smartAccount config", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => ({
+        smartAccount: {
+          chain: { id: 1, name: "test", rpc: "test" },
+        },
+      }),
+    } as Wallet;
+    expect(isSmartWallet(wallet)).toBe(true);
+  });
+
+  it("should return false for non-smart wallet", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => ({}),
+    } as Wallet;
+    expect(isSmartWallet(wallet)).toBe(false);
+  });
+
+  it("should return false when getConfig returns null", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => null,
+      // biome-ignore lint/suspicious/noExplicitAny: Testing invalid config
+    } as any as Wallet;
+    expect(isSmartWallet(wallet)).toBe(false);
+  });
+
+  it("should return false when getConfig returns undefined", () => {
+    const wallet = {
+      id: "inApp",
+      getConfig: () => undefined,
+    } as Wallet;
+    expect(isSmartWallet(wallet)).toBe(false);
+  });
+
+  it("should return false when smartAccount is null", () => {
+    const wallet = {
+      id: "inApp",
+      // biome-ignore lint/suspicious/noExplicitAny: Testing invalid config
+      getConfig: () => ({ smartAccount: null }) as any,
+    } as Wallet;
+    expect(isSmartWallet(wallet)).toBe(false);
+  });
+});

packages/thirdweb/src/wallets/smart/is-smart-wallet.ts

@@ -0,0 +1,21 @@
+import type { Wallet } from "../interfaces/wallet.js";
+
+/**
+ * Checks if the given wallet is a smart wallet.
+ *
+ * @param {Wallet} wallet - The wallet to check.
+ * @returns {boolean} True if the wallet is a smart wallet, false otherwise.
+ * @internal
+ */
+export function isSmartWallet(wallet: Wallet): boolean {
+  if (wallet.id === "smart") {
+    return true;
+  }
+
+  const config = wallet.getConfig();
+  if (!!config && "smartAccount" in config && !!config.smartAccount) {
+    return true;
+  }
+
+  return false;
+}

packages/thirdweb/src/wallets/smart/lib/smart-wallet-capabilities.ts

@@ -4,7 +4,7 @@ import type { Wallet } from "../../interfaces/wallet.js";
  * @internal
  */
 export function smartWalletGetCapabilities(args: {
-  wallet: Wallet<"smart">;
+  wallet: Wallet;
 }) {
   const { wallet } = args;