Add beforeSign hook

Author: matthewwalsh0

packages/transaction-controller/src/TransactionController.ts

@@ -117,6 +117,7 @@ import type {
   AfterAddHook,
   GasFeeEstimateLevel as GasFeeEstimateLevelType,
   TransactionBatchMeta,
+  AfterSimulateHook,
 } from './types';
 import {
   GasFeeEstimateLevel,
@@ -414,6 +415,9 @@ export type TransactionControllerOptions = {
       signedTx: TypedTransaction,
     ) => boolean;
 
+    /** Additional logic to execute after simulating a transaction. */
+    afterSimulate?: AfterSimulateHook;
+
     /**
      * Additional logic to execute before checking pending transactions.
      * Return false to prevent the broadcast of the transaction.
@@ -701,6 +705,8 @@ export class TransactionController extends BaseController<
     signedTx: TypedTransaction,
   ) => boolean;
 
+  readonly #afterSimulate: AfterSimulateHook;
+
   readonly #approvingTransactionIds: Set<string> = new Set();
 
   readonly #beforeCheckPendingTransaction: (
@@ -846,6 +852,7 @@ export class TransactionController extends BaseController<
 
     this.#afterAdd = hooks?.afterAdd ?? (() => Promise.resolve({}));
     this.#afterSign = hooks?.afterSign ?? (() => true);
+    this.#afterSimulate = hooks?.afterSimulate ?? (() => Promise.resolve({}));
     this.#beforeCheckPendingTransaction =
       /* istanbul ignore next */
       hooks?.beforeCheckPendingTransaction ?? (() => Promise.resolve(true));
@@ -1305,6 +1312,7 @@ export class TransactionController extends BaseController<
 
       if (requireApproval !== false) {
         this.#updateSimulationData(addedTransactionMeta, {
+          ethQuery,
           traceContext,
         }).catch((error) => {
           log('Error while updating simulation data', error);
@@ -4051,9 +4059,11 @@ export class TransactionController extends BaseController<
     transactionMeta: TransactionMeta,
     {
       blockTime,
+      ethQuery,
       traceContext,
     }: {
       blockTime?: number;
+      ethQuery?: EthQuery;
       traceContext?: TraceContext;
     } = {},
   ) {
@@ -4092,7 +4102,7 @@ export class TransactionController extends BaseController<
       }
     }
 
-    const finalTransactionMeta = this.#getTransaction(transactionId);
+    let finalTransactionMeta = this.#getTransaction(transactionId);
 
     /* istanbul ignore if */
     if (!finalTransactionMeta) {
@@ -4105,7 +4115,7 @@ export class TransactionController extends BaseController<
       return;
     }
 
-    this.#updateTransactionInternal(
+    finalTransactionMeta = this.#updateTransactionInternal(
       {
         transactionId,
         note: 'TransactionController#updateSimulationData - Update simulation data',
@@ -4117,7 +4127,39 @@ export class TransactionController extends BaseController<
       },
     );
 
-    log('Updated simulation data', transactionId, simulationData);
+    log('Calling afterSimulate hook', finalTransactionMeta);
+
+    if (ethQuery) {
+      const result = await this.#afterSimulate({
+        transactionMeta: finalTransactionMeta,
+      });
+
+      if (result.updateTransaction) {
+        log('Updating transaction with afterSimulate data');
+
+        finalTransactionMeta = this.#updateTransactionInternal(
+          { transactionId },
+          result.updateTransaction,
+        );
+
+        const { estimatedGas } = await estimateGas({
+          chainId: finalTransactionMeta.chainId,
+          ethQuery,
+          isSimulationEnabled: this.#isSimulationEnabled(),
+          messenger: this.messagingSystem,
+          txParams: { ...finalTransactionMeta.txParams, gas: undefined },
+        });
+
+        finalTransactionMeta = this.#updateTransactionInternal(
+          { transactionId },
+          (txMeta) => {
+            txMeta.txParams.gas = estimatedGas;
+          },
+        );
+      }
+    }
+
+    log('Updated simulation data', transactionId, finalTransactionMeta);
   }
 
   #onGasFeePollerTransactionUpdate({

packages/transaction-controller/src/index.ts

@@ -32,6 +32,7 @@ export {
 } from './TransactionController';
 export type {
   AfterAddHook,
+  AfterSimulateHook,
   Authorization,
   AuthorizationList,
   BatchTransactionParams,

packages/transaction-controller/src/types.ts

@@ -1817,3 +1817,13 @@ export type AfterAddHook = (request: {
 }) => Promise<{
   updateTransaction?: (transaction: TransactionMeta) => void;
 }>;
+
+/**
+ * Custom logic to be executed after a transaction is simulated.
+ * Can optionally update the transaction by returning the `updateTransaction` callback.
+ */
+export type AfterSimulateHook = (request: {
+  transactionMeta: TransactionMeta;
+}) => Promise<{
+  updateTransaction?: (transaction: TransactionMeta) => void;
+}>;