feat(relay-kit): Add UserOperation custom nonce support (#1126)

packages/relay-kit/src/packs/safe-4337/Safe4337Pack.test.ts

@@ -1042,5 +1042,35 @@ describe('Safe4337Pack', () => {
         '5afe006137303238633936636562316132623939353333646561393063346135'
       )
     })
+
+    it('should allow to use custom nonces', async () => {
+      const transferUSDC = {
+        to: fixtures.PAYMASTER_TOKEN_ADDRESS,
+        data: generateTransferCallData(fixtures.SAFE_ADDRESS_v1_4_1_WITH_0_3_0_MODULE, 100_000n),
+        value: '0',
+        operation: 0
+      }
+
+      const safe4337Pack = await createSafe4337Pack({
+        options: {
+          safeAddress: fixtures.SAFE_ADDRESS_v1_4_1_WITH_0_3_0_MODULE
+        }
+      })
+
+      const customNonce = utils.encodeNonce({
+        key: BigInt(Date.now()),
+        sequence: 0n
+      })
+
+      let safeOperation = await safe4337Pack.createTransaction({
+        transactions: [transferUSDC],
+        options: { customNonce }
+      })
+
+      expect(safeOperation.getUserOperation()).toHaveProperty('nonce', customNonce.toString())
+      expect(safeOperation.getSafeOperation()).toHaveProperty('nonce', customNonce.toString())
+
+      safeOperation = await safe4337Pack.signSafeOperation(safeOperation)
+    })
   })
 })

packages/relay-kit/src/packs/safe-4337/Safe4337Pack.ts

@@ -463,12 +463,13 @@ export class Safe4337Pack extends RelayKitBasePack<{
     transactions,
     options = {}
   }: Safe4337CreateTransactionProps): Promise<BaseSafeOperation> {
-    const { amountToApprove, validUntil, validAfter, feeEstimator } = options
+    const { amountToApprove, validUntil, validAfter, feeEstimator, customNonce } = options
 
     const userOperation = await createUserOperation(this.protocolKit, transactions, {
       entryPoint: this.#ENTRYPOINT_ADDRESS,
       paymasterOptions: this.#paymasterOptions,
-      amountToApprove
+      amountToApprove,
+      customNonce
     })
 
     if (this.#onchainIdentifier) {

packages/relay-kit/src/packs/safe-4337/types.ts

@@ -77,6 +77,7 @@ export type Safe4337CreateTransactionProps = {
     validUntil?: number
     validAfter?: number
     feeEstimator?: IFeeEstimator
+    customNonce?: bigint
   }
 }
 

packages/relay-kit/src/packs/safe-4337/utils/encodeNonce.ts

@@ -0,0 +1,8 @@
+import { toHex } from 'viem'
+
+export function encodeNonce(args: { key: bigint; sequence: bigint }): bigint {
+  const key = BigInt(toHex(args.key, { size: 24 }))
+  const sequence = BigInt(toHex(args.sequence, { size: 8 }))
+
+  return (key << BigInt(64)) + sequence
+}

packages/relay-kit/src/packs/safe-4337/utils/index.ts

@@ -47,3 +47,4 @@ export * from './entrypoint'
 export * from './signing'
 export * from './userOperations'
 export * from './getRelayKitVersion'
+export * from './encodeNonce'

packages/relay-kit/src/packs/safe-4337/utils/userOperations.ts

@@ -114,11 +114,18 @@ export async function createUserOperation(
   {
     amountToApprove,
     entryPoint,
-    paymasterOptions
-  }: { entryPoint: string; amountToApprove?: bigint; paymasterOptions: PaymasterOptions }
+    paymasterOptions,
+    customNonce
+  }: {
+    entryPoint: string
+    amountToApprove?: bigint
+    paymasterOptions: PaymasterOptions
+    customNonce?: bigint
+  }
 ): Promise<UserOperation> {
   const safeAddress = await protocolKit.getAddress()
-  const nonce = await getSafeNonceFromEntrypoint(protocolKit, safeAddress, entryPoint)
+  const nonce =
+    customNonce || (await getSafeNonceFromEntrypoint(protocolKit, safeAddress, entryPoint))
   const isSafeDeployed = await protocolKit.isSafeDeployed()
   const paymasterAndData =
     paymasterOptions && 'paymasterAddress' in paymasterOptions

playground/config/run.ts

@@ -25,7 +25,9 @@ const playgroundRelayKitPaths = {
   'userop-erc20-paymaster': 'relay-kit/userop-erc20-paymaster',
   'userop-erc20-paymaster-counterfactual': 'relay-kit/userop-erc20-paymaster-counterfactual',
   'userop-verifying-paymaster': 'relay-kit/userop-verifying-paymaster',
-  'userop-verifying-paymaster-counterfactual': 'relay-kit/userop-verifying-paymaster-counterfactual'
+  'userop-verifying-paymaster-counterfactual':
+    'relay-kit/userop-verifying-paymaster-counterfactual',
+  'userop-parallel-execution': 'relay-kit/userop-parallel-execution'
 }
 
 const playgroundStarterKitPaths = {

playground/relay-kit/userop-parallel-execution.ts

@@ -0,0 +1,81 @@
+import * as dotenv from 'dotenv'
+import { encodeNonce, Safe4337Pack } from '@safe-global/relay-kit'
+import { waitForOperationToFinish, setup4337Playground } from '../utils'
+
+dotenv.config({ path: './playground/relay-kit/.env' })
+
+const {
+  PRIVATE_KEY,
+  SAFE_ADDRESS = '0x',
+  RPC_URL = '',
+  CHAIN_ID = '',
+  BUNDLER_URL = ''
+} = process.env
+
+// PIM test token contract address
+// faucet: https://dashboard.pimlico.io/test-erc20-faucet
+const pimlicoTokenAddress = '0xFC3e86566895Fb007c6A0d3809eb2827DF94F751'
+
+const NUMBER_OF_OPERATIONS = 2
+
+async function main() {
+  // 1) Initialize pack
+  const safe4337Pack = await Safe4337Pack.init({
+    provider: RPC_URL,
+    signer: PRIVATE_KEY,
+    bundlerUrl: BUNDLER_URL,
+    safeModulesVersion: '0.3.0', // Blank or 0.3.0 for Entrypoint v0.7, 0.2.0 for Entrypoint v0.6
+    options: {
+      safeAddress: SAFE_ADDRESS
+    }
+  })
+
+  // 2) Setup Playground
+  const { transactions, timestamp } = await setup4337Playground(safe4337Pack, {
+    // nativeTokenAmount: parseEther('0.05'), // Increase this value when is not enough to cover the gas fees
+    erc20TokenAmount: 200_000n,
+    erc20TokenContractAddress: pimlicoTokenAddress
+  })
+
+  // 3) Create Multiple SafeOperations
+  const safeOperations = []
+
+  for (let i = 0; i < NUMBER_OF_OPERATIONS; i++) {
+    safeOperations.push(
+      safe4337Pack.createTransaction({
+        transactions,
+        options: {
+          validAfter: Number(timestamp - 60_000n),
+          validUntil: Number(timestamp + 60_000n),
+          customNonce: encodeNonce({
+            key: BigInt(Date.now()) + BigInt(i), // Ensure unique nonce
+            sequence: 0n
+          })
+        }
+      })
+    )
+  }
+
+  const createdSafeOperations = await Promise.all(safeOperations)
+
+  // 4) Sign all SafeOperations
+  const signingPromises = createdSafeOperations.map((op) => safe4337Pack.signSafeOperation(op))
+  const signedOperations = await Promise.all(signingPromises)
+
+  // Log all operations
+  signedOperations.forEach((op, index) => console.log(`SafeOperation ${index + 1}`, op))
+
+  // 5) Execute all operations in parallel
+  const executionPromises = signedOperations.map((op) =>
+    safe4337Pack.executeTransaction({ executable: op })
+  )
+
+  const userOperationHashes = await Promise.all(executionPromises)
+
+  // Wait for all operations to complete
+  await Promise.all(
+    userOperationHashes.map((hash) => waitForOperationToFinish(hash, CHAIN_ID, safe4337Pack))
+  )
+}
+
+main()