feat: add entity id and nonce reader

account-kit/smart-contracts/src/ma-v2/abis/entityIdAndNonceReader.ts

@@ -0,0 +1,23 @@
+export const entityIdAndNonceReaderAbi = [
+  {
+    type: "constructor",
+    inputs: [
+      {
+        name: "account",
+        type: "address",
+        internalType: "contract IModularAccountView",
+      },
+      {
+        name: "ep",
+        type: "address",
+        internalType: "contract IEntryPoint",
+      },
+      {
+        name: "nonce",
+        type: "uint192",
+        internalType: "uint192",
+      },
+    ],
+    stateMutability: "nonpayable",
+  },
+];

account-kit/smart-contracts/src/ma-v2/actions/DeferralActions.ts → account-kit/smart-contracts/src/ma-v2/actions/deferralActions.ts

@@ -15,7 +15,11 @@ import {
   encodePacked,
   size,
   toHex,
+  encodeDeployData,
+  hexToNumber,
 } from "viem";
+import { entityIdAndNonceReaderBytecode, buildFullNonceKey } from "../utils.js";
+import { entityIdAndNonceReaderAbi } from "../abis/entityIdAndNonceReader.js";
 import type { ModularAccountV2Client } from "../client/client.js";
 
 export type DeferredActionTypedData = {
@@ -62,6 +66,12 @@ export type BuildUserOperationWithDeferredActionParams = {
   nonceOverride: bigint;
 };
 
+export type EntityIdAndNonceParams = {
+  entityId: number;
+  nonceKey: bigint;
+  isGlobalValidation: boolean;
+};
+
 export type DeferralActions = {
   createDeferredActionTypedDataObject: (
     args: CreateDeferredActionTypedDataParams
@@ -70,6 +80,9 @@ export type DeferralActions = {
   buildUserOperationWithDeferredAction: (
     args: BuildUserOperationWithDeferredActionParams
   ) => Promise<UserOperationRequest_v7>;
+  getEntityIdAndNonce: (
+    args: EntityIdAndNonceParams
+  ) => Promise<{ nonce: bigint; entityId: number }>;
 };
 
 /**
@@ -86,15 +99,14 @@ export const deferralActions: (
     deadline,
     entityId,
     isGlobalValidation,
-    nonceKeyOverride,
+    nonceKeyOverride = 0n,
   }: CreateDeferredActionTypedDataParams): Promise<DeferredActionReturnData> => {
     if (!client.account) {
       throw new AccountNotFoundError();
     }
 
-    const baseNonceKey = nonceKeyOverride || 0n;
-    if (baseNonceKey > maxUint152) {
-      throw new InvalidNonceKeyError(baseNonceKey);
+    if (nonceKeyOverride > maxUint152) {
+      throw new InvalidNonceKeyError(nonceKeyOverride);
     }
 
     const entryPoint = client.account.getEntryPoint();
@@ -110,10 +122,12 @@ export const deferralActions: (
 
     // 2 = deferred action flags    0b10
     // 1 = isGlobal validation flag 0b01
-    const fullNonceKey: bigint =
-      ((baseNonceKey << 40n) + (BigInt(entityId) << 8n)) |
-      2n |
-      (isGlobalValidation ? 1n : 0n);
+    const fullNonceKey: bigint = buildFullNonceKey({
+      nonceKey: nonceKeyOverride,
+      entityId,
+      isGlobalValidation,
+      isDeferredAction: true,
+    });
 
     const nonceOverride = (await entryPointContract.read.getNonce([
       client.account.address,
@@ -224,9 +238,54 @@ export const deferralActions: (
     return unsignedUo;
   };
 
+  const getEntityIdAndNonce = async ({
+    entityId,
+    nonceKey,
+    isGlobalValidation,
+  }: EntityIdAndNonceParams) => {
+    if (!client.account) {
+      throw new AccountNotFoundError();
+    }
+
+    if (nonceKey > maxUint152) {
+      throw new InvalidNonceKeyError(nonceKey);
+    }
+
+    const entryPoint = client.account.getEntryPoint();
+    if (entryPoint === undefined) {
+      throw new EntryPointNotFoundError(client.chain, "0.7.0");
+    }
+
+    const bytecode = encodeDeployData({
+      abi: entityIdAndNonceReaderAbi,
+      bytecode: entityIdAndNonceReaderBytecode,
+      args: [
+        client.account.address,
+        entryPoint.address,
+        buildFullNonceKey({
+          nonceKey,
+          entityId,
+          isGlobalValidation,
+          isDeferredAction: true,
+        }),
+      ],
+    });
+
+    const { data } = await client.call({ data: bytecode });
+    if (!data) {
+      throw new Error("No data returned from contract call");
+    }
+
+    return {
+      nonce: BigInt(data),
+      entityId: hexToNumber(`0x${data.slice(40, 48)}`),
+    };
+  };
+
   return {
     createDeferredActionTypedDataObject,
     buildDeferredActionDigest,
     buildUserOperationWithDeferredAction,
+    getEntityIdAndNonce,
   };
 };

account-kit/smart-contracts/src/ma-v2/client/client.test.ts

@@ -16,12 +16,14 @@ import {
   hashMessage,
   hashTypedData,
   fromHex,
+  prepareEncodeFunctionData,
   isAddress,
   concat,
   testActions,
-  type TestActions,
   concatHex,
+  type TestActions,
   type Hex,
+  type ContractFunctionName,
 } from "viem";
 import { HookType } from "../actions/common/types.js";
 import {
@@ -37,6 +39,7 @@ import {
   AllowlistModule,
   NativeTokenLimitModule,
   semiModularAccountBytecodeAbi,
+  buildFullNonceKey,
 } from "@account-kit/smart-contracts/experimental";
 import {
   createLightAccountClient,
@@ -58,7 +61,7 @@ import {
   alchemyGasAndPaymasterAndDataMiddleware,
 } from "@account-kit/infra";
 import { getMAV2UpgradeToData } from "@account-kit/smart-contracts";
-import { deferralActions } from "../actions/DeferralActions.js";
+import { deferralActions } from "../actions/deferralActions.js";
 import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
 
 // Note: These tests maintain a shared state to not break the local-running rundler by desyncing the chain.
@@ -1345,6 +1348,129 @@ describe("MA v2 Tests", async () => {
     client.setAutomine(true);
   });
 
+  it("tests entity id and nonce selection", async () => {
+    let newClient = (await givenConnectedProvider({ signer, salt: 1n }))
+      .extend(deferralActions)
+      .extend(installValidationActions);
+
+    await setBalance(client, {
+      address: newClient.getAddress(),
+      value: parseEther("2"),
+    });
+
+    const entryPoint = newClient.account.getEntryPoint();
+    const entryPointContract = getContract({
+      address: entryPoint.address,
+      abi: entryPoint.abi,
+      client,
+    });
+
+    // entity id and nonce selection for undeployed account
+    for (let startEntityId = 0; startEntityId < 5; startEntityId++) {
+      for (let startNonce = 0n; startNonce < 5n; startNonce++) {
+        const { entityId, nonce } = await newClient.getEntityIdAndNonce({
+          entityId: startEntityId,
+          nonceKey: startNonce,
+          isGlobalValidation: true,
+        });
+
+        const expectedEntityId: number = Math.max(1, startEntityId);
+
+        // account not deployed, expect to get 1 when we pass in 0
+        expect(entityId).toEqual(expectedEntityId);
+        await expect(
+          entryPointContract.read.getNonce([
+            newClient.account.address,
+            buildFullNonceKey({
+              nonceKey: startNonce,
+              entityId: expectedEntityId,
+              isDeferredAction: true,
+            }),
+          ])
+        ).resolves.toEqual(nonce);
+      }
+    }
+
+    // deploy the account and install at entity id 1 with global validation
+    const uo1 = await newClient.installValidation({
+      validationConfig: {
+        moduleAddress: getDefaultSingleSignerValidationModuleAddress(
+          newClient.chain
+        ),
+        entityId: 1,
+        isGlobal: true,
+        isSignatureValidation: false,
+        isUserOpValidation: false,
+      },
+      selectors: [],
+      installData: SingleSignerValidationModule.encodeOnInstallData({
+        entityId: 1,
+        signer: await sessionKey.getAddress(),
+      }),
+      hooks: [],
+    });
+    await newClient.waitForUserOperationTransaction(uo1);
+
+    const fns: ContractFunctionName<typeof semiModularAccountBytecodeAbi>[] = [
+      "execute",
+      "executeBatch",
+    ];
+
+    const selectors = fns.map(
+      (s) =>
+        prepareEncodeFunctionData({
+          abi: semiModularAccountBytecodeAbi,
+          functionName: s,
+        }).functionName
+    );
+
+    // deploy the account and install some entity ids with selector validation
+    const uo2 = await newClient.installValidation({
+      validationConfig: {
+        moduleAddress: getDefaultSingleSignerValidationModuleAddress(
+          newClient.chain
+        ),
+        entityId: 2,
+        isGlobal: false,
+        isSignatureValidation: false,
+        isUserOpValidation: false,
+      },
+      selectors,
+      installData: SingleSignerValidationModule.encodeOnInstallData({
+        entityId: 2,
+        signer: await sessionKey.getAddress(),
+      }),
+      hooks: [],
+    });
+    await newClient.waitForUserOperationTransaction(uo2);
+
+    // entity id and nonce selection for undeployed account
+    for (let startEntityId = 1; startEntityId < 5; startEntityId++) {
+      for (let startNonce = 0n; startNonce < 5n; startNonce++) {
+        const { entityId, nonce } = await newClient.getEntityIdAndNonce({
+          entityId: startEntityId,
+          nonceKey: startNonce,
+          isGlobalValidation: true,
+        });
+
+        const expectedEntityId: number = Math.max(startEntityId, 3);
+
+        // expect to get max(3, startEntityId)
+        expect(entityId).toEqual(expectedEntityId);
+        await expect(
+          entryPointContract.read.getNonce([
+            newClient.account.address,
+            buildFullNonceKey({
+              nonceKey: startNonce,
+              entityId: expectedEntityId,
+              isDeferredAction: true,
+            }),
+          ])
+        ).resolves.toEqual(nonce);
+      }
+    }
+  });
+
   it("upgrade from a lightaccount", async () => {
     const lightAccountClient = await createLightAccountClient({
       chain: instance.chain,
@@ -1398,11 +1524,13 @@ describe("MA v2 Tests", async () => {
     signerEntity,
     accountAddress,
     paymasterMiddleware,
+    salt = 0n,
   }: {
     signer: SmartAccountSigner;
     signerEntity?: SignerEntity;
     accountAddress?: `0x${string}`;
     paymasterMiddleware?: "alchemyGasAndPaymasterAndData" | "erc7677";
+    salt?: bigint;
   }) =>
     createModularAccountV2Client({
       chain: instance.chain,
@@ -1419,6 +1547,7 @@ describe("MA v2 Tests", async () => {
         : paymasterMiddleware === "erc7677"
         ? erc7677Middleware()
         : {}),
+      salt,
     });
 
   it("alchemy client calls the createAlchemySmartAccountClient", async () => {

account-kit/smart-contracts/src/ma-v2/index.ts

@@ -29,6 +29,7 @@ export {
   getDefaultTimeRangeModuleAddress,
   getDefaultWebauthnValidationModuleAddress,
 } from "./modules/utils.js";
+export { buildFullNonceKey } from "./utils.js";
 export { allowlistModuleAbi } from "./modules/allowlist-module/abis/allowlistModuleAbi.js";
 export { AllowlistModule } from "./modules/allowlist-module/module.js";
 export { nativeTokenLimitModuleAbi } from "./modules/native-token-limit-module/abis/nativeTokenLimitModuleAbi.js";

account-kit/smart-contracts/src/ma-v2/utils.ts

@@ -228,3 +228,28 @@ export async function getMAV2UpgradeToData<
       }),
   };
 }
+
+export const entityIdAndNonceReaderBytecode =
+  "0x608060405234801561001057600080fd5b506040516104f13803806104f183398101604081905261002f916101e5565b60006008826001600160c01b0316901c90506000808263ffffffff1611610057576001610059565b815b90506001600160a01b0385163b15610133575b60006001600160a01b03861663d31b575b6bffffffff0000000000000000604085901b166040516001600160e01b031960e084901b1681526001600160401b03199091166004820152602401600060405180830381865afa1580156100d5573d6000803e3d6000fd5b505050506040513d6000823e601f3d908101601f191682016040526100fd91908101906103c6565b805190915060ff161580156101155750606081015151155b156101205750610133565b8161012a816104a4565b9250505061006c565b604051631aab3f0d60e11b81526001600160a01b03868116600483015264ffffffff01600160c01b038516600884901b64ffffffff0016176024830152600091908616906335567e1a90604401602060405180830381865afa15801561019d573d6000803e3d6000fd5b505050506040513d601f19601f820116820180604052508101906101c191906104d7565b90508060005260206000f35b6001600160a01b03811681146101e257600080fd5b50565b6000806000606084860312156101fa57600080fd5b8351610205816101cd565b6020850151909350610216816101cd565b60408501519092506001600160c01b038116811461023357600080fd5b809150509250925092565b634e487b7160e01b600052604160045260246000fd5b604051608081016001600160401b03811182821017156102765761027661023e565b60405290565b604051601f8201601f191681016001600160401b03811182821017156102a4576102a461023e565b604052919050565b60006001600160401b038211156102c5576102c561023e565b5060051b60200190565b600082601f8301126102e057600080fd5b81516102f36102ee826102ac565b61027c565b8082825260208201915060208360051b86010192508583111561031557600080fd5b602085015b8381101561034857805166ffffffffffffff198116811461033a57600080fd5b83526020928301920161031a565b5095945050505050565b600082601f83011261036357600080fd5b81516103716102ee826102ac565b8082825260208201915060208360051b86010192508583111561039357600080fd5b602085015b838110156103485780516001600160e01b0319811681146103b857600080fd5b835260209283019201610398565b6000602082840312156103d857600080fd5b81516001600160401b038111156103ee57600080fd5b82016080818503121561040057600080fd5b610408610254565b815160ff8116811461041957600080fd5b815260208201516001600160401b0381111561043457600080fd5b610440868285016102cf565b60208301525060408201516001600160401b0381111561045f57600080fd5b61046b868285016102cf565b60408301525060608201516001600160401b0381111561048a57600080fd5b61049686828501610352565b606083015250949350505050565b600063ffffffff821663ffffffff81036104ce57634e487b7160e01b600052601160045260246000fd5b60010192915050565b6000602082840312156104e957600080fd5b505191905056fe";
+
+export type BuildNonceParams = {
+  nonceKey?: bigint;
+  entityId?: number;
+  isGlobalValidation?: boolean;
+  isDeferredAction?: boolean;
+  isDirectCallValidation?: boolean;
+};
+
+export const buildFullNonceKey = ({
+  nonceKey = 0n,
+  entityId = 0,
+  isGlobalValidation = true,
+  isDeferredAction = false,
+}: BuildNonceParams): bigint => {
+  return (
+    (nonceKey << 40n) +
+    BigInt(entityId << 8) +
+    (isDeferredAction ? 2n : 0n) +
+    (isGlobalValidation ? 1n : 0n)
+  );
+};