refactor: change how the middleware is defined

packages/core/src/middleware/alchemy-paymaster.ts

@@ -2,6 +2,7 @@ import type { Address, Hex } from "viem";
 import type { PublicErc4337Client } from "../client/types.js";
 import type { UserOperationRequest, UserOperationStruct } from "../types.js";
 import { deepHexlify, resolveProperties } from "../utils.js";
+import type { ISmartAccountProvider } from "../provider/types.js";
 
 type ClientWithAlchemyMethod = PublicErc4337Client & {
   request: PublicErc4337Client["request"] &
@@ -27,13 +28,14 @@ export interface AlchemyPaymasterConfig {
 
 export const alchemyPaymasterAndDataMiddleware = (
   config: AlchemyPaymasterConfig
-) => ({
-  dummyPaymasterMiddleware: async (struct: UserOperationStruct) => {
-    struct.paymasterAndData =
-      "0xc03aac639bb21233e0139381970328db8bceeb67fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c";
-    return struct;
+): Parameters<ISmartAccountProvider["withPaymasterMiddleware"]>["0"] => ({
+  dummyPaymasterDataMiddleware: async (_struct: UserOperationStruct) => {
+    return {
+      paymasterAndData:
+        "0xc03aac639bb21233e0139381970328db8bceeb67fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c",
+    };
   },
-  getPaymasterAndDataMiddleware: async (struct: UserOperationStruct) => {
+  paymasterDataMiddleware: async (struct: UserOperationStruct) => {
     const { paymasterAndData } = await (
       config.provider as ClientWithAlchemyMethod
     ).request({
@@ -47,7 +49,6 @@ export const alchemyPaymasterAndDataMiddleware = (
       ],
     });
 
-    struct.paymasterAndData = paymasterAndData;
-    return struct;
+    return { paymasterAndData };
   },
 });

packages/core/src/provider/base.ts

@@ -21,7 +21,11 @@ import {
 } from "../utils.js";
 import type {
   AccountMiddlewareFn,
+  AccountMiddlewareOverrideFn,
+  FeeDataMiddleware,
+  GasEstimatorMiddleware,
   ISmartAccountProvider,
+  PaymasterAndDataMiddleware,
   SendUserOperationResult,
 } from "./types.js";
 
@@ -174,7 +178,7 @@ export class SmartAccountProvider<
       this.dummyPaymasterDataMiddleware,
       this.gasEstimator,
       this.feeDataGetter,
-      this.paymasterMiddleware
+      this.paymasterDataMiddleware
     )({
       initCode,
       sender: this.getAddress(),
@@ -222,11 +226,10 @@ export class SmartAccountProvider<
     return struct;
   };
 
-  readonly paymasterMiddleware: AccountMiddlewareFn = async (
+  readonly paymasterDataMiddleware: AccountMiddlewareFn = async (
     struct: UserOperationStruct
   ): Promise<UserOperationStruct> => {
     struct.paymasterAndData = "0x";
-
     return struct;
   };
 
@@ -271,30 +274,37 @@ export class SmartAccountProvider<
   };
 
   withPaymasterMiddleware = (overrides: {
-    dummyPaymasterMiddleware?: AccountMiddlewareFn;
-    getPaymasterAndDataMiddleware?: AccountMiddlewareFn;
+    dummyPaymasterDataMiddleware?: PaymasterAndDataMiddleware;
+    paymasterDataMiddleware?: PaymasterAndDataMiddleware;
   }): this => {
-    defineReadOnly(
-      this,
-      "dummyPaymasterDataMiddleware",
-      overrides.dummyPaymasterMiddleware ?? this.dummyPaymasterDataMiddleware
-    );
-    defineReadOnly(
-      this,
-      "paymasterMiddleware",
-      overrides.getPaymasterAndDataMiddleware ?? this.paymasterMiddleware
-    );
+    const newDummyMiddleware = overrides.dummyPaymasterDataMiddleware
+      ? this.overrideMiddlewareFunction(overrides.dummyPaymasterDataMiddleware)
+      : this.dummyPaymasterDataMiddleware;
+    defineReadOnly(this, "dummyPaymasterDataMiddleware", newDummyMiddleware);
+
+    const newPaymasterMiddleware = overrides.paymasterDataMiddleware
+      ? this.overrideMiddlewareFunction(overrides.paymasterDataMiddleware)
+      : this.paymasterDataMiddleware;
+    defineReadOnly(this, "paymasterDataMiddleware", newPaymasterMiddleware);
 
     return this;
   };
 
-  withGasEstimator = (override: AccountMiddlewareFn): this => {
-    defineReadOnly(this, "gasEstimator", override);
+  withGasEstimator = (override: GasEstimatorMiddleware): this => {
+    defineReadOnly(
+      this,
+      "gasEstimator",
+      this.overrideMiddlewareFunction(override)
+    );
     return this;
   };
 
-  withFeeDataGetter = (override: AccountMiddlewareFn): this => {
-    defineReadOnly(this, "feeDataGetter", override);
+  withFeeDataGetter = (override: FeeDataMiddleware): this => {
+    defineReadOnly(
+      this,
+      "feeDataGetter",
+      this.overrideMiddlewareFunction(override)
+    );
     return this;
   };
 
@@ -305,4 +315,15 @@ export class SmartAccountProvider<
     defineReadOnly(this, "account", account);
     return this as this & { account: typeof account };
   }
+
+  private overrideMiddlewareFunction = (
+    override: AccountMiddlewareOverrideFn
+  ): AccountMiddlewareFn => {
+    return async (struct: UserOperationStruct) => {
+      return {
+        ...struct,
+        ...(await override(struct)),
+      };
+    };
+  };
 }

packages/core/src/provider/types.ts

@@ -16,13 +16,29 @@ export type AccountMiddlewareFn = (
   struct: UserOperationStruct
 ) => Promise<UserOperationStruct>;
 
+export type AccountMiddlewareOverrideFn<
+  K extends keyof UserOperationStruct = never
+> = (
+  struct: UserOperationStruct
+) => Promise<Required<Pick<UserOperationStruct, K>>>;
+
+export type PaymasterAndDataMiddleware =
+  AccountMiddlewareOverrideFn<"paymasterAndData">;
+
+export type GasEstimatorMiddleware = AccountMiddlewareOverrideFn<
+  "callGasLimit" | "preVerificationGas" | "verificationGasLimit"
+>;
+export type FeeDataMiddleware = AccountMiddlewareOverrideFn<
+  "maxFeePerGas" | "maxPriorityFeePerGas"
+>;
+
 // TODO: this also will need to implement EventEmitteer
 export interface ISmartAccountProvider<
   TTransport extends SupportedTransports = Transport
 > {
   readonly rpcClient: PublicErc4337Client<TTransport>;
   readonly dummyPaymasterDataMiddleware: AccountMiddlewareFn;
-  readonly paymasterMiddleware: AccountMiddlewareFn;
+  readonly paymasterDataMiddleware: AccountMiddlewareFn;
   readonly gasEstimator: AccountMiddlewareFn;
   readonly feeDataGetter: AccountMiddlewareFn;
 
@@ -85,8 +101,8 @@ export interface ISmartAccountProvider<
    * @returns an update instance of this, which now uses the new middleware
    */
   withPaymasterMiddleware: (overrides: {
-    dummyPaymasterMiddleware?: AccountMiddlewareFn;
-    getPaymasterAndDataMiddleware?: AccountMiddlewareFn;
+    dummyPaymasterDataMiddleware?: PaymasterAndDataMiddleware;
+    paymasterDataMiddleware?: PaymasterAndDataMiddleware;
   }) => this;
 
   /**
@@ -96,7 +112,7 @@ export interface ISmartAccountProvider<
    * @param override - a function for overriding the default gas estimator middleware
    * @returns
    */
-  withGasEstimator: (override: AccountMiddlewareFn) => this;
+  withGasEstimator: (override: GasEstimatorMiddleware) => this;
 
   /**
    * Overrides the feeDataGetter middleware which is used for setting the fee fields on the UserOperation
@@ -105,7 +121,7 @@ export interface ISmartAccountProvider<
    * @param override - a function for overriding the default feeDataGetter middleware
    * @returns
    */
-  withFeeDataGetter: (override: AccountMiddlewareFn) => this;
+  withFeeDataGetter: (override: FeeDataMiddleware) => this;
 
   /**
    * Sets the current account to the account returned by the given function. The function parameter

packages/core/src/utils.ts

@@ -32,7 +32,7 @@ export const asyncPipe =
     for (const fn of fns) {
       result = await fn(result);
     }
-    return x;
+    return result;
   };
 
 // based on @ethersproject/properties, but pulled in here to minize the dependency on ethers