feat: optimize Transaction Pay amount quote latency

packages/transaction-controller/CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `updateTransactionCallback` for atomic callback-based metadata updates and export `updateEIP7702BatchData` for synchronous indexed updates to nested transaction calldata ([#9543](https://github.com/MetaMask/core/pull/9543))
+
+### Fixed
+
+- Clear stale gas and simulation metadata synchronously when updating EIP-7702 batch calldata ([#9543](https://github.com/MetaMask/core/pull/9543))
+
 ## [69.0.0]
 
 ### Changed

packages/transaction-controller/src/TransactionController-method-action-types.ts

@@ -134,6 +134,18 @@ export type TransactionControllerUpdateTransactionAction = {
   handler: TransactionController['updateTransaction'];
 };
 
+/**
+ * Updates an existing transaction using a callback.
+ *
+ * @param transactionId - ID of the transaction to update.
+ * @param callback - Function that updates the transaction metadata.
+ * @returns The updated transaction metadata.
+ */
+export type TransactionControllerUpdateTransactionCallbackAction = {
+  type: `TransactionController:updateTransactionCallback`;
+  handler: TransactionController['updateTransactionCallback'];
+};
+
 /**
  * Mark a transaction as failed, transitioning it through the standard failure
  * path.
@@ -441,6 +453,7 @@ export type TransactionControllerMethodActions =
   | TransactionControllerEstimateGasBatchAction
   | TransactionControllerEstimateGasBufferedAction
   | TransactionControllerUpdateTransactionAction
+  | TransactionControllerUpdateTransactionCallbackAction
   | TransactionControllerFailTransactionAction
   | TransactionControllerUpdateSecurityAlertResponseAction
   | TransactionControllerWipeTransactionsAction

packages/transaction-controller/src/TransactionController.test.ts

@@ -4925,6 +4925,50 @@ describe('TransactionController', () => {
     });
   });
 
+  describe('updateTransactionCallback', () => {
+    it('updates multiple properties using a callback and returns the updated transaction', () => {
+      const { controller } = setupController({
+        options: {
+          state: {
+            transactions: [TRANSACTION_META_MOCK],
+          },
+        },
+      });
+
+      const result = controller.updateTransactionCallback(
+        TRANSACTION_META_MOCK.id,
+        (transactionMeta) => {
+          transactionMeta.requiredAssets = [
+            {
+              address: ACCOUNT_2_MOCK,
+              amount: '0x1',
+              standard: 'erc20',
+            },
+          ];
+          transactionMeta.txParams.value = '0x2';
+        },
+      );
+
+      expect(result).toStrictEqual(controller.state.transactions[0]);
+      expect(result.requiredAssets).toStrictEqual([
+        {
+          address: ACCOUNT_2_MOCK,
+          amount: '0x1',
+          standard: 'erc20',
+        },
+      ]);
+      expect(result.txParams.value).toBe('0x2');
+    });
+
+    it('throws if the transaction does not exist', () => {
+      const { controller } = setupController();
+
+      expect(() =>
+        controller.updateTransactionCallback('missing-id', () => undefined),
+      ).toThrow('Cannot update transaction as ID not found - missing-id');
+    });
+  });
+
   describe('updateTransactionGasFees', () => {
     it('throws if transaction does not exist', async () => {
       const { controller } = setupController();
@@ -7683,6 +7727,75 @@ describe('TransactionController', () => {
       expect(result).not.toContain('4567');
     });
 
+    it('clears stale preparation metadata before gas estimation completes', async () => {
+      const gasPreparation = createDeferredPromise<void>();
+      const receiptRevert = { message: 'Receipt reverted' };
+      const simulationRevert = { message: 'Simulation reverted' };
+      updateGasMock.mockImplementationOnce(async ({ txMeta }) => {
+        await gasPreparation.promise;
+        txMeta.txParams.gas = '0x222';
+        txMeta.gasLimitNoBuffer = '0x200';
+      });
+      const { controller } = setupController({
+        options: {
+          state: {
+            transactions: [
+              {
+                ...TRANSACTION_META_MOCK,
+                gasLimitNoBuffer: '0x100',
+                gasUsed: '0x101',
+                nestedTransactions: [{ to: ACCOUNT_2_MOCK, data: '0x1234' }],
+                revert: {
+                  gas: { message: 'Gas reverted' },
+                  receipt: receiptRevert,
+                  simulation: simulationRevert,
+                },
+                securityAlertResponse: {
+                  reason: 'Previous revision warning',
+                  result_type: 'Warning',
+                },
+                simulationData: SIMULATION_DATA_RESULT_MOCK,
+                simulationFails: {
+                  debug: {},
+                  reason: 'Previous gas estimate failed',
+                },
+                txParams: {
+                  ...TRANSACTION_META_MOCK.txParams,
+                  gas: '0x102',
+                },
+              },
+            ],
+          },
+        },
+      });
+
+      const updatePromise = controller.updateAtomicBatchData({
+        transactionId: TRANSACTION_META_MOCK.id,
+        transactionIndex: 0,
+        transactionData: '0x89AB',
+      });
+      const transaction = controller.state.transactions[0];
+
+      expect(transaction.nestedTransactions?.[0].data).toBe('0x89AB');
+      expect(transaction.txParams.data).toContain('89ab');
+      expect(transaction.txParams.gas).toBeUndefined();
+      expect(transaction.gasLimitNoBuffer).toBeUndefined();
+      expect(transaction.gasUsed).toBeUndefined();
+      expect(transaction.securityAlertResponse).toBeUndefined();
+      expect(transaction.simulationData).toBeUndefined();
+      expect(transaction.simulationFails).toBeUndefined();
+      expect(transaction.revert).toStrictEqual({
+        receipt: receiptRevert,
+        simulation: simulationRevert,
+      });
+
+      gasPreparation.resolve();
+      await updatePromise;
+
+      expect(controller.state.transactions[0].txParams.gas).toBe('0x222');
+      expect(controller.state.transactions[0].gasLimitNoBuffer).toBe('0x200');
+    });
+
     it('updates gas', async () => {
       const gasMock = '0x1234';
       const gasLimitNoBufferMock = '0x123';
@@ -8374,6 +8487,29 @@ describe('TransactionController', () => {
       });
     });
 
+    describe('TransactionController:updateTransactionCallback', () => {
+      it('calls updateTransactionCallback via messenger', () => {
+        const { controller, messenger } = setupController({
+          options: {
+            state: {
+              transactions: [TRANSACTION_META_MOCK],
+            },
+          },
+        });
+
+        const result = messenger.call(
+          'TransactionController:updateTransactionCallback',
+          TRANSACTION_META_MOCK.id,
+          (transactionMeta) => {
+            transactionMeta.txParams.value = '0x1';
+          },
+        );
+
+        expect(result).toStrictEqual(controller.state.transactions[0]);
+        expect(result.txParams.value).toBe('0x1');
+      });
+    });
+
     describe('TransactionController:getGasFeeTokens', () => {
       it('returns gas fee tokens', async () => {
         const { messenger } = setupController();

packages/transaction-controller/src/TransactionController.ts

@@ -141,9 +141,9 @@ import {
 import { getBalanceChanges } from './utils/balance-changes';
 import { addTransactionBatch, isAtomicBatchSupported } from './utils/batch';
 import {
-  generateEIP7702BatchTransaction,
   getDelegationAddress,
   signAuthorizationList,
+  updateEIP7702BatchData,
 } from './utils/eip7702';
 import { validateConfirmedExternalTransaction } from './utils/external-transactions';
 import {
@@ -697,6 +697,7 @@ const MESSENGER_EXPOSED_METHODS = [
   'updateSecurityAlertResponse',
   'updateSelectedGasFeeToken',
   'updateTransaction',
+  'updateTransactionCallback',
   'updateTransactionGasFees',
   'wipeTransactions',
 ] as const;
@@ -1618,6 +1619,20 @@ export class TransactionController extends BaseController<
     log('Transaction updated', { transactionId, note });
   }
 
+  /**
+   * Updates an existing transaction using a callback.
+   *
+   * @param transactionId - ID of the transaction to update.
+   * @param callback - Function that updates the transaction metadata.
+   * @returns The updated transaction metadata.
+   */
+  updateTransactionCallback(
+    transactionId: string,
+    callback: (transactionMeta: TransactionMeta) => TransactionMeta | void,
+  ): Readonly<TransactionMeta> {
+    return this.#updateTransactionInternal({ transactionId }, callback);
+  }
+
   /**
    * Mark a transaction as failed, transitioning it through the standard failure
    * path.
@@ -2576,63 +2591,70 @@ export class TransactionController extends BaseController<
       transactionData,
     });
 
+    const currentTransaction = this.#getTransaction(transactionId);
+
+    if (!currentTransaction) {
+      throw new Error(
+        `Cannot update transaction as ID not found - ${transactionId}`,
+      );
+    }
+
+    const { nestedTransactions, transactionData: updatedTransactionData } =
+      updateEIP7702BatchData(
+        currentTransaction.txParams.from as Hex,
+        currentTransaction.nestedTransactions ?? [],
+        [{ transactionIndex, transactionData }],
+      );
     const updatedTransactionMeta = this.#updateTransactionInternal(
-      {
-        transactionId,
-      },
+      { transactionId },
       (transactionMeta) => {
-        const { nestedTransactions, txParams } = transactionMeta;
-        const from = txParams.from as Hex;
-        const nestedTransaction = nestedTransactions?.[transactionIndex];
+        transactionMeta.nestedTransactions = nestedTransactions;
+        transactionMeta.txParams.data = updatedTransactionData;
+        transactionMeta.txParams.gas = undefined;
+        transactionMeta.gasLimitNoBuffer = undefined;
+        transactionMeta.gasUsed = undefined;
+        transactionMeta.securityAlertResponse = undefined;
+        transactionMeta.simulationData = undefined;
+        transactionMeta.simulationFails = undefined;
+
+        if (transactionMeta.revert) {
+          delete transactionMeta.revert.gas;
 
-        if (!nestedTransaction) {
-          throw new Error(
-            `Nested transaction not found with index - ${transactionIndex}`,
-          );
+          if (
+            !transactionMeta.revert.simulation &&
+            !transactionMeta.revert.receipt
+          ) {
+            transactionMeta.revert = undefined;
+          }
         }
-
-        nestedTransaction.data = transactionData;
-
-        const batchTransaction = generateEIP7702BatchTransaction(
-          from,
-          nestedTransactions,
-        );
-
-        transactionMeta.txParams.data = batchTransaction.data;
       },
     );
-
     const draftTransaction = cloneDeep({
       ...updatedTransactionMeta,
       txParams: {
         ...updatedTransactionMeta.txParams,
-        // Clear existing gas to force estimation
+        // Clear existing gas to force estimation.
         gas: undefined,
       },
     });
 
     await this.#updateGasEstimate(draftTransaction);
 
-    this.#updateTransactionInternal(
-      {
-        transactionId,
-      },
-      (transactionMeta) => {
-        transactionMeta.txParams.gas = draftTransaction.txParams.gas;
-        transactionMeta.simulationFails = draftTransaction.simulationFails;
-        transactionMeta.gasLimitNoBuffer = draftTransaction.gasLimitNoBuffer;
-
-        const draftGasRevert = draftTransaction.revert?.gas;
-        if (draftGasRevert) {
-          transactionMeta.revert = {
-            ...transactionMeta.revert,
-            gas: draftGasRevert,
-          };
-        }
-      },
-    );
+    this.#updateTransactionInternal({ transactionId }, (transactionMeta) => {
+      transactionMeta.txParams.gas = draftTransaction.txParams.gas;
+      transactionMeta.simulationFails = draftTransaction.simulationFails;
+      transactionMeta.gasLimitNoBuffer = draftTransaction.gasLimitNoBuffer;
+
+      const draftGasRevert = draftTransaction.revert?.gas;
+      if (draftGasRevert) {
+        transactionMeta.revert = {
+          ...transactionMeta.revert,
+          gas: draftGasRevert,
+        };
+      }
+    });
 
-    return updatedTransactionMeta.txParams.data as Hex;
+    return updatedTransactionData;
   }
 
   /**

packages/transaction-controller/src/index.ts

@@ -36,6 +36,7 @@ export type {
   TransactionControllerGetTransactionsAction,
   TransactionControllerUpdateCustodialTransactionAction,
   TransactionControllerUpdateTransactionAction,
+  TransactionControllerUpdateTransactionCallbackAction,
   TransactionControllerHandleMethodDataAction,
   TransactionControllerIsAtomicBatchSupportedAction,
   TransactionControllerStopTransactionAction,
@@ -89,6 +90,7 @@ export type {
   Log,
   MetamaskPayMetadata,
   NestedTransactionMetadata,
+  NestedTransactionUpdate,
   PublishBatchHook,
   PublishBatchHookRequest,
   PublishBatchHookResult,
@@ -132,6 +134,7 @@ export { mergeGasFeeEstimates } from './utils/gas-flow';
 export {
   decodeAuthorizationSignature,
   generateEIP7702BatchTransaction,
+  updateEIP7702BatchData,
 } from './utils/eip7702';
 export {
   isEIP1559Transaction,

packages/transaction-controller/src/types.ts

@@ -2331,6 +2331,15 @@ export type RequiredAsset = {
   standard: string;
 };
 
+/** A nested transaction calldata update in an atomic batch. */
+export type NestedTransactionUpdate = {
+  /** Index of the nested transaction to update. */
+  transactionIndex: number;
+
+  /** New calldata for the nested transaction. */
+  transactionData: Hex;
+};
+
 /**
  * Decoded revert from a single lifecycle source.
  */