feat: optimize Transaction Pay amount quote latency

packages/transaction-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `beginAtomicBatchUpdate` for coherent multi-call updates with monotonic revisions and revision-bound gas preparation ([#9543](https://github.com/MetaMask/core/pull/9543))
+
 ## [69.0.0]
 
 ### Changed

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

@@ -358,6 +358,18 @@ export type TransactionControllerAbortTransactionSigningAction = {
   handler: TransactionController['abortTransactionSigning'];
 };
 
+/**
+ * Atomically updates all amount-dependent data in an atomic batch and starts
+ * revision-bound local preparation.
+ *
+ * @param request - Complete atomic batch update.
+ * @returns The synchronously published revision and its preparation handle.
+ */
+export type TransactionControllerBeginAtomicBatchUpdateAction = {
+  type: `TransactionController:beginAtomicBatchUpdate`;
+  handler: TransactionController['beginAtomicBatchUpdate'];
+};
+
 /**
  * Update the transaction data of a single nested transaction within an atomic batch transaction.
  *
@@ -457,6 +469,7 @@ export type TransactionControllerMethodActions =
   | TransactionControllerGetLayer1GasFeeAction
   | TransactionControllerClearUnapprovedTransactionsAction
   | TransactionControllerAbortTransactionSigningAction
+  | TransactionControllerBeginAtomicBatchUpdateAction
   | TransactionControllerUpdateAtomicBatchDataAction
   | TransactionControllerUpdateSelectedGasFeeTokenAction
   | TransactionControllerUpdateRequiredTransactionIdsAction

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

@@ -7618,6 +7618,247 @@ describe('TransactionController', () => {
     });
   });
 
+  describe('beginAtomicBatchUpdate', () => {
+    const requiredAssets = [
+      {
+        address: ACCOUNT_2_MOCK,
+        amount: '0x64' as Hex,
+        standard: 'erc20',
+      },
+    ];
+
+    function setupAtomicBatchController(
+      transactionOverrides: Partial<TransactionMeta> = {},
+    ): ReturnType<typeof setupController> {
+      return setupController({
+        options: {
+          state: {
+            transactions: [
+              {
+                ...TRANSACTION_META_MOCK,
+                status: TransactionStatus.unapproved,
+                nestedTransactions: [
+                  { to: ACCOUNT_2_MOCK, data: '0x1234' },
+                  { to: ACCOUNT_2_MOCK, data: '0x4567' },
+                ],
+                ...transactionOverrides,
+              },
+            ],
+          },
+        },
+        updateToInitialState: true,
+      });
+    }
+
+    it('publishes all nested updates and required assets synchronously in one revision', async () => {
+      const { controller } = setupAtomicBatchController();
+
+      const result = controller.beginAtomicBatchUpdate({
+        transactionId: TRANSACTION_META_MOCK.id,
+        requiredAssets,
+        nestedTransactionUpdates: [
+          { transactionIndex: 0, transactionData: '0xAAAA' },
+          { transactionIndex: 1, transactionData: '0xBBBB' },
+        ],
+      });
+
+      expect(result).toStrictEqual({
+        revision: 1,
+        transaction: expect.objectContaining({
+          requiredAssets,
+          transactionRevision: 1,
+        }),
+        preparation: expect.any(Promise),
+      });
+      expect(
+        controller.state.transactions[0].nestedTransactions?.map(
+          ({ data }) => data,
+        ),
+      ).toStrictEqual(['0xAAAA', '0xBBBB']);
+      expect(controller.state.transactions[0].txParams.data).toContain('aaaa');
+      expect(controller.state.transactions[0].txParams.data).toContain('bbbb');
+
+      expect(await result.preparation).toMatchObject({
+        revision: 1,
+        status: 'prepared',
+      });
+      expect(updateGasMock).toHaveBeenCalledTimes(1);
+    });
+
+    it.each([
+      {
+        expectedError: 'at least one nested transaction update',
+        name: 'an empty patch',
+        updates: [],
+      },
+      {
+        expectedError: 'Duplicate nested transaction index - 0',
+        name: 'duplicate indexes',
+        updates: [
+          { transactionIndex: 0, transactionData: '0xAAAA' as Hex },
+          { transactionIndex: 0, transactionData: '0xBBBB' as Hex },
+        ],
+      },
+      {
+        expectedError: 'Nested transaction not found with index - 2',
+        name: 'an invalid index',
+        updates: [{ transactionIndex: 2, transactionData: '0xAAAA' as Hex }],
+      },
+    ])(
+      'rejects $name without publishing state',
+      ({ expectedError, updates }) => {
+        const { controller } = setupAtomicBatchController();
+        const stateBeforeUpdate = controller.state;
+
+        expect(() =>
+          controller.beginAtomicBatchUpdate({
+            transactionId: TRANSACTION_META_MOCK.id,
+            requiredAssets,
+            nestedTransactionUpdates: updates,
+          }),
+        ).toThrow(expectedError);
+        expect(controller.state).toBe(stateBeforeUpdate);
+        expect(updateGasMock).not.toHaveBeenCalled();
+      },
+    );
+
+    it('clears stale revision-bound preparation metadata while preserving non-gas reverts', async () => {
+      const gasPreparation = createDeferredPromise<void>();
+      const simulationRevert = { message: 'Simulation reverted' };
+      const receiptRevert = { message: 'Receipt reverted' };
+      updateGasMock.mockImplementationOnce(async () => {
+        await gasPreparation.promise;
+      });
+      const { controller } = setupAtomicBatchController({
+        gasLimitNoBuffer: '0x111',
+        gasUsed: '0x222',
+        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: '0x333',
+        },
+      });
+
+      const result = controller.beginAtomicBatchUpdate({
+        transactionId: TRANSACTION_META_MOCK.id,
+        requiredAssets,
+        nestedTransactionUpdates: [
+          { transactionIndex: 0, transactionData: '0xAAAA' },
+        ],
+      });
+      const transaction = controller.state.transactions[0];
+
+      expect(transaction.txParams.gas).toBeUndefined();
+      expect(transaction.gasLimitNoBuffer).toBeUndefined();
+      expect(transaction.gasUsed).toBeUndefined();
+      expect(transaction.simulationData).toBeUndefined();
+      expect(transaction.simulationFails).toBeUndefined();
+      expect(transaction.securityAlertResponse).toBeUndefined();
+      expect(transaction.revert).toStrictEqual({
+        receipt: receiptRevert,
+        simulation: simulationRevert,
+      });
+
+      gasPreparation.resolve();
+      await result.preparation;
+    });
+
+    it('ignores simulation results started before the current revision', async () => {
+      const staleSimulation = createDeferredPromise<{
+        simulationData: SimulationData;
+      }>();
+      const { controller } = setupAtomicBatchController();
+      getBalanceChangesMock.mockReturnValueOnce(staleSimulation.promise);
+      shouldResimulateMock.mockReturnValueOnce({
+        blockTime: 123,
+        resimulate: true,
+      });
+
+      await controller.updateEditableParams(TRANSACTION_META_MOCK.id, {});
+      expect(getBalanceChangesMock).toHaveBeenCalledTimes(1);
+
+      const update = controller.beginAtomicBatchUpdate({
+        transactionId: TRANSACTION_META_MOCK.id,
+        requiredAssets,
+        nestedTransactionUpdates: [
+          { transactionIndex: 0, transactionData: '0xAAAA' },
+        ],
+      });
+      await update.preparation;
+
+      staleSimulation.resolve({
+        simulationData: {
+          ...SIMULATION_DATA_RESULT_MOCK,
+          nativeBalanceChange: undefined,
+        },
+      });
+      await flushPromises();
+
+      expect(controller.state.transactions[0].simulationData).toStrictEqual(
+        SIMULATION_DATA_RESULT_MOCK,
+      );
+      expect(controller.state.transactions[0].transactionRevision).toBe(1);
+    });
+
+    it('increments revisions and prevents stale gas from overwriting a newer revision', async () => {
+      const firstGas = createDeferredPromise<void>();
+      const { controller } = setupAtomicBatchController();
+
+      updateGasMock
+        .mockImplementationOnce(async ({ txMeta }) => {
+          await firstGas.promise;
+          txMeta.txParams.gas = '0x111';
+        })
+        .mockImplementationOnce(async ({ txMeta }) => {
+          txMeta.txParams.gas = '0x222';
+        });
+
+      const first = controller.beginAtomicBatchUpdate({
+        transactionId: TRANSACTION_META_MOCK.id,
+        requiredAssets,
+        nestedTransactionUpdates: [
+          { transactionIndex: 0, transactionData: '0xAAAA' },
+        ],
+      });
+      const second = controller.beginAtomicBatchUpdate({
+        transactionId: TRANSACTION_META_MOCK.id,
+        requiredAssets,
+        nestedTransactionUpdates: [
+          { transactionIndex: 0, transactionData: '0xBBBB' },
+        ],
+      });
+
+      expect(await second.preparation).toMatchObject({
+        revision: 2,
+        status: 'prepared',
+      });
+      firstGas.resolve();
+      expect(await first.preparation).toMatchObject({
+        revision: 1,
+        status: 'superseded',
+      });
+
+      expect(controller.state.transactions[0].transactionRevision).toBe(2);
+      expect(controller.state.transactions[0].txParams.gas).toBe('0x222');
+      expect(
+        controller.state.transactions[0].nestedTransactions?.[0].data,
+      ).toBe('0xBBBB');
+    });
+  });
+
   describe('updateAtomicBatchData', () => {
     /**
      * Template for updateAtomicBatchData test.