feat: add upgrade functionality for light account to msca

packages/accounts/src/light-account/account.ts

@@ -1,25 +1,35 @@
 import {
   SimpleSmartContractAccount,
-  SmartAccountProvider,
+  type ISmartAccountProvider,
   type SignTypedDataParams,
   type SmartAccountSigner,
 } from "@alchemy/aa-core";
 import {
   concatHex,
   decodeFunctionResult,
   encodeFunctionData,
+  fromHex,
+  trim,
   type Address,
+  type Chain,
   type FallbackTransport,
   type Hash,
   type Hex,
   type Transport,
 } from "viem";
+import { createMultiOwnerMSCA } from "../msca/multi-owner-account.js";
+import { getDefaultMSCAFactoryAddress } from "../msca/utils.js";
 import { LightAccountAbi } from "./abis/LightAccountAbi.js";
 import { LightAccountFactoryAbi } from "./abis/LightAccountFactoryAbi.js";
 
 export class LightSmartContractAccount<
   TTransport extends Transport | FallbackTransport = Transport
 > extends SimpleSmartContractAccount<TTransport> {
+  static readonly implementationAddress =
+    "0x5467b1947f47d0646704eb801e075e72aeae8113";
+  static readonly storageSlot =
+    "0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
+
   override async signTypedData(params: SignTypedDataParams): Promise<Hash> {
     return this.owner.signTypedData(params);
   }
@@ -55,6 +65,92 @@ export class LightSmartContractAccount<
     return decodedCallResult;
   }
 
+  /**
+   * Upgrades the account implementation from Light Account to a Modular Account.
+   * Optionally waits for the transaction to be mined.
+   *
+   * @param provider - the provider to use to send the transaction
+   * @param chain - the chain to upgrade the account on
+   * @param accountImplAddress - the address of the smart account implementation to
+   * upgrade to
+   * @param initializationData - the initialization data address to use when upgrading to the new
+   * smart account
+   * @param waitForTxn - whether or not to wait for the transaction to be mined
+   * @returns {
+   *  provider: SmartAccountProvider<TTransport> & { account: MSCA };
+   *  hash: Hash;
+   * } - the upgraded provider and corresponding userOperation hash,
+   * or transaction hash if `waitForTxn` is true
+   */
+  static async upgrade<
+    P extends ISmartAccountProvider,
+    TTransport extends Transport | FallbackTransport = Transport
+  >(
+    provider: P & {
+      account: LightSmartContractAccount<TTransport>;
+    },
+    chain: Chain,
+    accountImplAddress: Address,
+    initializationData: Hex,
+    waitForTxn: boolean = false
+  ) {
+    const accountAddress = await provider.getAddress();
+
+    const storage = await provider.rpcClient.getStorageAt({
+      address: accountAddress,
+      slot: LightSmartContractAccount.storageSlot,
+    });
+
+    if (storage == null) {
+      throw new Error("could not get storage");
+    }
+
+    // only upgrade undeployed accounts (storage 0) or deployed light accounts, error otherwise
+    if (
+      fromHex(storage, "number") !== 0 &&
+      trim(storage) !== LightSmartContractAccount.implementationAddress
+    ) {
+      throw new Error(
+        "could not determine if smart account implementation is light account"
+      );
+    }
+
+    const encodeUpgradeData = encodeFunctionData({
+      abi: LightAccountAbi,
+      functionName: "upgradeToAndCall",
+      args: [accountImplAddress, initializationData],
+    });
+
+    const result = await provider.sendUserOperation({
+      target: accountAddress,
+      data: encodeUpgradeData,
+    });
+
+    let hash = result.hash;
+    if (waitForTxn) {
+      hash = await provider.waitForUserOperationTransaction(result.hash);
+    }
+
+    const owner = provider.account.getOwner();
+    if (owner == null) {
+      throw new Error("could not get owner");
+    }
+
+    return {
+      provider: provider.connect((rpcClient) =>
+        createMultiOwnerMSCA({
+          rpcClient,
+          factoryAddress: getDefaultMSCAFactoryAddress(chain),
+          owner,
+          index: 0n,
+          chain: chain,
+          accountAddress,
+        })
+      ),
+      hash,
+    };
+  }
+
   /**
    * Encodes the transferOwnership function call using Light Account ABI.
    *
@@ -79,9 +175,10 @@ export class LightSmartContractAccount<
    * @returns {Hash} the userOperation hash, or transaction hash if `waitForTxn` is true
    */
   static async transferOwnership<
+    P extends ISmartAccountProvider,
     TTransport extends Transport | FallbackTransport = Transport
   >(
-    provider: SmartAccountProvider<TTransport> & {
+    provider: P & {
       account: LightSmartContractAccount<TTransport>;
     },
     newOwner: SmartAccountSigner,

packages/accounts/src/light-account/e2e-tests/light-account.test.ts

@@ -19,13 +19,18 @@ import {
   createLightAccountProvider,
   LightSmartContractAccount,
 } from "../../index.js";
+import {
+  defaultMSCAImplementationAddress,
+  getMSCAInitializationData,
+} from "../../msca/utils.js";
 import {
   API_KEY,
   LIGHT_ACCOUNT_OWNER_MNEMONIC,
   UNDEPLOYED_OWNER_MNEMONIC,
 } from "./constants.js";
 
 const chain = sepolia;
+const multiOwnerPluginAddress = "0x56bC629F342821FBe91C5273880792dFECBE7920";
 
 Logger.setLogLevel(LogLevel.DEBUG);
 
@@ -195,6 +200,52 @@ describe("Light Account Tests", () => {
     expect(newOwnerViaProvider).not.toBe(oldOwner);
     expect(newOwnerViaProvider).toBe(newOwner);
   }, 100000);
+
+  it("should upgrade a deployed light account to msca successfully", async () => {
+    const provider = givenConnectedProvider({
+      owner,
+      chain,
+    });
+
+    // create a throwaway address
+    const throwawayOwner = LocalAccountSigner.privateKeyToAccountSigner(
+      generatePrivateKey()
+    );
+    const throwawayProvider = givenConnectedProvider({
+      owner: throwawayOwner,
+      chain,
+    });
+
+    const accountAddress = await throwawayProvider.getAddress();
+    const ownerAddress = await throwawayOwner.getAddress();
+
+    // fund + deploy the throwaway address
+    await provider.sendTransaction({
+      from: await provider.getAddress(),
+      to: accountAddress,
+      data: "0x",
+      value: toHex(1000000000000000n),
+    });
+
+    const { provider: upgradedProvider } =
+      await LightSmartContractAccount.upgrade(
+        throwawayProvider,
+        chain,
+        defaultMSCAImplementationAddress,
+        await getMSCAInitializationData(
+          throwawayProvider,
+          multiOwnerPluginAddress
+        ),
+        true
+      );
+
+    const upgradedAccountAddress = await upgradedProvider.getAddress();
+
+    const owners = await upgradedProvider.account.readOwners();
+
+    expect(upgradedAccountAddress).toBe(accountAddress);
+    expect(owners).toContain(ownerAddress);
+  }, 200000);
 });
 
 const givenConnectedProvider = ({