refactor: optimize gas estimates with overrides

Author: moldy530

aa-sdk/core/src/index.ts

@@ -23,6 +23,7 @@ export { sendTransactions } from "./actions/smartAccount/sendTransactions.js";
 export { sendUserOperation } from "./actions/smartAccount/sendUserOperation.js";
 export type * from "./actions/smartAccount/types.js";
 export { waitForUserOperationTransaction } from "./actions/smartAccount/waitForUserOperationTransacation.js";
+export * from "./client/addBreadcrumb.js";
 export type * from "./client/bundlerClient.js";
 export {
   createBundlerClient,
@@ -65,19 +66,19 @@ export {
   IncorrectAccountType,
   SignTransactionNotSupportedError,
   SmartAccountWithSignerRequiredError,
-  UpgradeToAndCallNotSupportedError,
   UpgradesNotSupportedError,
+  UpgradeToAndCallNotSupportedError,
 } from "./errors/account.js";
 export { BaseError } from "./errors/base.js";
 export {
   ChainNotFoundError,
+  EntityIdOverrideError,
   IncompatibleClientError,
-  InvalidRpcUrlError,
+  InvalidDeferredActionNonce,
   InvalidEntityIdError,
-  InvalidNonceKeyError,
-  EntityIdOverrideError,
   InvalidModularAccountV2Mode,
-  InvalidDeferredActionNonce,
+  InvalidNonceKeyError,
+  InvalidRpcUrlError,
 } from "./errors/client.js";
 export {
   EntryPointNotFoundError,
@@ -92,11 +93,10 @@ export {
   InvalidUserOperationError,
   WaitForUserOperationError,
 } from "./errors/useroperation.js";
-export * from "./client/addBreadcrumb.js";
-export { LogLevel, Logger } from "./logger.js";
+export { Logger, LogLevel } from "./logger.js";
 export { middlewareActions } from "./middleware/actions.js";
-export { default7702UserOpSigner } from "./middleware/defaults/7702signer.js";
 export { default7702GasEstimator } from "./middleware/defaults/7702gasEstimator.js";
+export { default7702UserOpSigner } from "./middleware/defaults/7702signer.js";
 export { defaultFeeEstimator } from "./middleware/defaults/feeEstimator.js";
 export { defaultGasEstimator } from "./middleware/defaults/gasEstimator.js";
 export { defaultPaymasterAndData } from "./middleware/defaults/paymasterAndData.js";
@@ -106,7 +106,7 @@ export { erc7677Middleware } from "./middleware/erc7677middleware.js";
 export { noopMiddleware } from "./middleware/noopMiddleware.js";
 export type * from "./middleware/types.js";
 export { LocalAccountSigner } from "./signer/local-account.js";
-export { SignerSchema, isSigner } from "./signer/schema.js";
+export { isSigner, SignerSchema } from "./signer/schema.js";
 export type {
   SmartAccountAuthenticator,
   SmartAccountSigner,
@@ -117,35 +117,36 @@ export { split, type SplitTransportParams } from "./transport/split.js";
 export type * from "./types.js";
 export type * from "./utils/index.js";
 export {
-  TraceHeader,
-  TRACE_HEADER_NAME,
-  TRACE_HEADER_STATE,
-} from "./utils/traceHeader.js";
-export {
-  BigNumberishRangeSchema,
-  BigNumberishSchema,
-  ChainSchema,
-  HexSchema,
-  MultiplierSchema,
   allEqual,
+  applyGasAndFeeOverrides,
   applyUserOpFeeOption,
   applyUserOpOverride,
   applyUserOpOverrideOrFeeOption,
   asyncPipe,
   bigIntMax,
   bigIntMultiply,
+  BigNumberishRangeSchema,
+  BigNumberishSchema,
   bypassPaymasterAndData,
   bypassPaymasterAndDataEmptyHex,
+  ChainSchema,
   concatPaymasterAndData,
   deepHexlify,
   filterUndefined,
   getDefaultUserOperationFeeOptions,
+  HexSchema,
   isBigNumberish,
   isMultiplier,
   isValidRequest,
+  MultiplierSchema,
   parsePaymasterAndData,
   pick,
   resolveProperties,
   takeBytes,
   toRecord,
 } from "./utils/index.js";
+export {
+  TRACE_HEADER_NAME,
+  TRACE_HEADER_STATE,
+  TraceHeader,
+} from "./utils/traceHeader.js";

aa-sdk/core/src/middleware/defaults/gasEstimator.ts

@@ -4,7 +4,10 @@ import type {
   UserOperationStruct,
 } from "../../types.js";
 import { deepHexlify, resolveProperties } from "../../utils/index.js";
-import { applyUserOpOverrideOrFeeOption } from "../../utils/userop.js";
+import {
+  applyGasAndFeeOverrides,
+  applyUserOpOverrideOrFeeOption,
+} from "../../utils/userop.js";
 import type { MiddlewareClient } from "../actions.js";
 import type { ClientMiddlewareFn } from "../types.js";
 
@@ -21,7 +24,9 @@ export const defaultGasEstimator: <C extends MiddlewareClient>(
 ) => ClientMiddlewareFn =
   (client: MiddlewareClient): ClientMiddlewareFn =>
   async (struct, { account, overrides, feeOptions }) => {
-    const request = deepHexlify(await resolveProperties(struct));
+    const request = deepHexlify(
+      await resolveProperties(await applyGasAndFeeOverrides(struct, overrides))
+    );
 
     const estimates = await client.estimateUserOperationGas(
       request,

aa-sdk/core/src/utils/userop.ts

@@ -1,4 +1,4 @@
-import { concat, type Address, type Hex } from "viem";
+import { concat, type Address, type Hex, type IsUndefined } from "viem";
 import type { EntryPointVersion } from "../entrypoint/types";
 import type {
   BigNumberish,
@@ -11,7 +11,12 @@ import type {
   UserOperationStruct_v7,
 } from "../types";
 import { bigIntClamp, bigIntMultiply } from "./bigint.js";
-import { allEqual, isBigNumberish } from "./index.js";
+import {
+  allEqual,
+  isBigNumberish,
+  resolveProperties,
+  type Deferrable,
+} from "./index.js";
 
 /**
  * Utility method for asserting a UserOperationStruct has valid fields for the given entry point version
@@ -105,6 +110,52 @@ export function applyUserOpOverride<TValue extends BigNumberish | undefined>(
   }
 }
 
+function keysOf<
+  T extends Record<string, unknown>,
+  Filter extends (keyof T)[] | undefined = undefined
+>(
+  obj?: T,
+  filterBy?: Filter
+): IsUndefined<Filter> extends true ? (keyof T)[] : NonNullable<Filter> {
+  if (obj == null) {
+    return [] as any;
+  }
+
+  const keys = Object.keys(obj);
+  if (filterBy == null) {
+    return keys as any;
+  }
+
+  return keys.filter((key) => filterBy.includes(key)) as any;
+}
+
+export async function applyGasAndFeeOverrides<
+  TEntryPointVersion extends EntryPointVersion
+>(
+  uo: Deferrable<UserOperationStruct<TEntryPointVersion>>,
+  overrides?: UserOperationOverrides<TEntryPointVersion>
+): Promise<Deferrable<UserOperationStruct<TEntryPointVersion>>> {
+  const awaitedUo = await resolveProperties(uo);
+
+  for (const override of keysOf(overrides, [
+    "callGasLimit",
+    "maxFeePerGas",
+    "maxPriorityFeePerGas",
+    "preVerificationGas",
+    "verificationGasLimit",
+  ])) {
+    const overrideValue = applyUserOpOverride(
+      awaitedUo[override],
+      overrides?.[override]
+    );
+    if (overrideValue != null) {
+      awaitedUo[override] = overrideValue;
+    }
+  }
+
+  return awaitedUo;
+}
+
 /**
  * Utility method for applying a UserOperationFeeOptionsField value
  * over the current value set for the field

account-kit/infra/src/middleware/gasManager.ts

@@ -11,6 +11,7 @@ import type {
   UserOperationRequest,
 } from "@aa-sdk/core";
 import {
+  applyGasAndFeeOverrides,
   bypassPaymasterAndData,
   ChainNotFoundError,
   clientHeaderTrack,
@@ -24,23 +25,23 @@ import {
   resolveProperties,
 } from "@aa-sdk/core";
 import {
-  fromHex,
-  isHex,
-  toHex,
-  type Hex,
   encodeAbiParameters,
   encodeFunctionData,
-  parseAbi,
+  fromHex,
+  isHex,
   maxUint256,
+  parseAbi,
   sliceHex,
+  toHex,
+  type Hex,
 } from "viem";
-import type { AlchemySmartAccountClient } from "../client/smartAccountClient.js";
+import type { MiddlewareClient } from "../../../../aa-sdk/core/dist/types/middleware/actions.js";
+import type { RequestGasAndPaymasterAndDataRequest } from "../actions/types.js";
 import type { AlchemyTransport } from "../alchemyTransport.js";
+import type { AlchemySmartAccountClient } from "../client/smartAccountClient.js";
+import type { PermitDomain, PermitMessage } from "../gas-manager.js";
+import { EIP712NoncesAbi, PermitTypes } from "../gas-manager.js";
 import { alchemyFeeEstimator } from "./feeEstimator.js";
-import type { RequestGasAndPaymasterAndDataRequest } from "../actions/types.js";
-import { PermitTypes, EIP712NoncesAbi } from "../gas-manager.js";
-import type { PermitMessage, PermitDomain } from "../gas-manager.js";
-import type { MiddlewareClient } from "../../../../aa-sdk/core/dist/types/middleware/actions.js";
 
 type Context = {
   policyId: string | string[];
@@ -119,6 +120,7 @@ export function alchemyGasManagerMiddleware(
     paymasterAndData: async (uo, args) => {
       const context = await buildContext(uo, args);
       const baseMiddleware = erc7677Middleware({ context });
+      uo = await applyGasAndFeeOverrides(uo, args.overrides);
       return baseMiddleware.paymasterAndData(uo, args);
     },
   };
@@ -223,39 +225,11 @@ export function alchemyGasAndPaymasterAndDataMiddleware(
         throw new ChainNotFoundError();
       }
 
-      const userOp = deepHexlify(await resolveProperties(uo));
+      const userOp = deepHexlify(
+        await resolveProperties(await applyGasAndFeeOverrides(uo, overrides_))
+      );
 
       const overrides: UserOperationOverrides = filterUndefined({
-        maxFeePerGas: overrideField(
-          "maxFeePerGas",
-          overrides_ as UserOperationOverrides,
-          feeOptions,
-          userOp
-        ),
-        maxPriorityFeePerGas: overrideField(
-          "maxPriorityFeePerGas",
-          overrides_ as UserOperationOverrides,
-          feeOptions,
-          userOp
-        ),
-        callGasLimit: overrideField(
-          "callGasLimit",
-          overrides_ as UserOperationOverrides,
-          feeOptions,
-          userOp
-        ),
-        verificationGasLimit: overrideField(
-          "verificationGasLimit",
-          overrides_ as UserOperationOverrides,
-          feeOptions,
-          userOp
-        ),
-        preVerificationGas: overrideField(
-          "preVerificationGas",
-          overrides_ as UserOperationOverrides,
-          feeOptions,
-          userOp
-        ),
         ...(account.getEntryPoint().version === "0.7.0"
           ? {
               paymasterVerificationGasLimit: overrideField<"0.7.0">(