feat: Add afterAdd hook to transaction controller

packages/transaction-controller/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add optional `afterAdd` hook to constructor ([#5692](https://github.com/MetaMask/core/pull/5692))
+  - Add optional `txParamsOriginal` property to `TransactionMeta`.
+  - Add `AfterAddHook` type.
+
 ## [54.1.0]
 
 ### Changed

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

@@ -2072,6 +2072,98 @@ describe('TransactionController', () => {
       );
     });
 
+    describe('with afterAdd hook', () => {
+      it('calls afterAdd hook', async () => {
+        const afterAddHook = jest.fn().mockResolvedValueOnce({});
+
+        const { controller } = setupController({
+          options: {
+            hooks: {
+              afterAdd: afterAddHook,
+            },
+          },
+        });
+
+        await controller.addTransaction(
+          {
+            from: ACCOUNT_MOCK,
+            to: ACCOUNT_MOCK,
+          },
+          {
+            networkClientId: NETWORK_CLIENT_ID_MOCK,
+          },
+        );
+
+        expect(afterAddHook).toHaveBeenCalledTimes(1);
+      });
+
+      it('updates transaction if update callback returned', async () => {
+        const updateTransactionMock = jest.fn();
+
+        const afterAddHook = jest
+          .fn()
+          .mockResolvedValueOnce({ updateTransaction: updateTransactionMock });
+
+        const { controller } = setupController({
+          options: {
+            hooks: {
+              afterAdd: afterAddHook,
+            },
+          },
+        });
+
+        await controller.addTransaction(
+          {
+            from: ACCOUNT_MOCK,
+            to: ACCOUNT_MOCK,
+          },
+          {
+            networkClientId: NETWORK_CLIENT_ID_MOCK,
+          },
+        );
+
+        expect(updateTransactionMock).toHaveBeenCalledTimes(1);
+        expect(updateTransactionMock).toHaveBeenCalledWith(
+          expect.objectContaining({
+            id: expect.any(String),
+          }),
+        );
+      });
+
+      it('saves original transaction params if update callback returned', async () => {
+        const updateTransactionMock = jest.fn();
+
+        const afterAddHook = jest
+          .fn()
+          .mockResolvedValueOnce({ updateTransaction: updateTransactionMock });
+
+        const { controller } = setupController({
+          options: {
+            hooks: {
+              afterAdd: afterAddHook,
+            },
+          },
+        });
+
+        await controller.addTransaction(
+          {
+            from: ACCOUNT_MOCK,
+            to: ACCOUNT_MOCK,
+          },
+          {
+            networkClientId: NETWORK_CLIENT_ID_MOCK,
+          },
+        );
+
+        expect(controller.state.transactions[0].txParamsOriginal).toStrictEqual(
+          expect.objectContaining({
+            from: ACCOUNT_MOCK,
+            to: ACCOUNT_MOCK,
+          }),
+        );
+      });
+    });
+
     describe('updates simulation data', () => {
       it('by default', async () => {
         getSimulationDataMock.mockResolvedValueOnce(

packages/transaction-controller/src/TransactionController.ts

@@ -110,6 +110,7 @@ import type {
   GasFeeToken,
   IsAtomicBatchSupportedResult,
   IsAtomicBatchSupportedRequest,
+  AfterAddHook,
 } from './types';
 import {
   TransactionEnvelopeType,
@@ -377,6 +378,9 @@ export type TransactionControllerOptions = {
 
   /** The controller hooks. */
   hooks: {
+    /** Additional logic to execute after adding a transaction. */
+    afterAdd?: AfterAddHook;
+
     /** Additional logic to execute after signing a transaction. Return false to not change the status to signed. */
     afterSign?: (
       transactionMeta: TransactionMeta,
@@ -662,6 +666,8 @@ export class TransactionController extends BaseController<
   TransactionControllerState,
   TransactionControllerMessenger
 > {
+  readonly #afterAdd: AfterAddHook;
+
   readonly #internalEvents = new EventEmitter();
 
   private readonly isHistoryDisabled: boolean;
@@ -891,6 +897,7 @@ export class TransactionController extends BaseController<
     this.#testGasFeeFlows = testGasFeeFlows === true;
     this.#trace = trace ?? (((_request, fn) => fn?.()) as TraceCallback);
 
+    this.#afterAdd = hooks?.afterAdd ?? (() => Promise.resolve({}));
     this.afterSign = hooks?.afterSign ?? (() => true);
     this.beforeCheckPendingTransaction =
       /* istanbul ignore next */
@@ -1247,6 +1254,20 @@ export class TransactionController extends BaseController<
           verifiedOnBlockchain: false,
         };
 
+    const { updateTransaction } = await this.#afterAdd({
+      transactionMeta: addedTransactionMeta,
+    });
+
+    if (updateTransaction) {
+      log('Updating transaction using afterAdd hook');
+
+      addedTransactionMeta.txParamsOriginal = cloneDeep(
+        addedTransactionMeta.txParams,
+      );
+
+      updateTransaction(addedTransactionMeta);
+    }
+
     await this.#trace(
       { name: 'Estimate Gas Properties', parentContext: traceContext },
       (context) =>

packages/transaction-controller/src/index.ts

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

packages/transaction-controller/src/types.ts

@@ -446,6 +446,11 @@ export type TransactionMeta = {
    */
   txParams: TransactionParams;
 
+  /**
+   * Initial transaction parameters before `afterAdd` hook was invoked.
+   */
+  txParamsOriginal?: TransactionParams;
+
   /**
    * Transaction receipt.
    */
@@ -1756,3 +1761,13 @@ export type IsAtomicBatchSupportedResultEntry = {
   /** Address of the contract that the account would be upgraded to. */
   upgradeContractAddress?: Hex;
 };
+
+/**
+ * Custom logic to be executed after a transaction is added.
+ * Can optionally update the transaction by returning the `updateTransaction` callback.
+ */
+export type AfterAddHook = (request: {
+  transactionMeta: TransactionMeta;
+}) => Promise<{
+  updateTransaction?: (transaction: TransactionMeta) => void;
+}>;