feat: Add tracing to the STX controller (#515)

Author: dan437

src/SmartTransactionsController.test.ts

@@ -22,7 +22,11 @@ import * as sinon from 'sinon';
 
 import packageJson from '../package.json';
 import { advanceTime, flushPromises, getFakeProvider } from '../tests/helpers';
-import { API_BASE_URL, SENTINEL_API_BASE_URL_MAP } from './constants';
+import {
+  API_BASE_URL,
+  SENTINEL_API_BASE_URL_MAP,
+  SmartTransactionsTraceName,
+} from './constants';
 import SmartTransactionsController, {
   DEFAULT_INTERVAL,
   getDefaultSmartTransactionsControllerState,
@@ -2387,6 +2391,151 @@ describe('SmartTransactionsController', () => {
       );
     });
   });
+
+  describe('Tracing', () => {
+    const createTraceCallback = () =>
+      jest.fn().mockImplementation(async (_request, fn) => {
+        return fn?.();
+      });
+
+    it('traces getFees API call with expected name', async () => {
+      const traceCallback = createTraceCallback();
+
+      await withController(
+        {
+          options: {
+            trace: traceCallback,
+          },
+        },
+        async ({ controller }) => {
+          const apiUrl = API_BASE_URL;
+          nock(apiUrl)
+            .post(`/networks/${ethereumChainIdDec}/getFees`)
+            .reply(200, createGetFeesApiResponse());
+
+          const tradeTx = createUnsignedTransaction(ethereumChainIdDec);
+          await controller.getFees(tradeTx);
+
+          expect(traceCallback).toHaveBeenCalledWith(
+            { name: SmartTransactionsTraceName.GetFees },
+            expect.any(Function),
+          );
+        },
+      );
+    });
+
+    it('traces submitSignedTransactions API call with expected name', async () => {
+      const traceCallback = createTraceCallback();
+
+      await withController(
+        {
+          options: {
+            trace: traceCallback,
+          },
+        },
+        async ({ controller }) => {
+          const apiUrl = API_BASE_URL;
+          nock(apiUrl)
+            .post(
+              `/networks/${ethereumChainIdDec}/submitTransactions?stxControllerVersion=${packageJson.version}`,
+            )
+            .reply(200, createSubmitTransactionsApiResponse());
+
+          const signedTx = createSignedTransaction();
+          const signedCanceledTx = createSignedCanceledTransaction();
+          const txParams = createTxParams();
+
+          await controller.submitSignedTransactions({
+            signedTransactions: [signedTx],
+            signedCanceledTransactions: [signedCanceledTx],
+            txParams,
+          });
+
+          expect(traceCallback).toHaveBeenCalledWith(
+            { name: SmartTransactionsTraceName.SubmitTransactions },
+            expect.any(Function),
+          );
+        },
+      );
+    });
+
+    it('traces cancelSmartTransaction API call with expected name', async () => {
+      const traceCallback = createTraceCallback();
+
+      await withController(
+        {
+          options: {
+            trace: traceCallback,
+          },
+        },
+        async ({ controller }) => {
+          const apiUrl = API_BASE_URL;
+          nock(apiUrl)
+            .post(`/networks/${ethereumChainIdDec}/cancel`)
+            .reply(200, {});
+
+          await controller.cancelSmartTransaction('uuid1');
+
+          expect(traceCallback).toHaveBeenCalledWith(
+            { name: SmartTransactionsTraceName.CancelTransaction },
+            expect.any(Function),
+          );
+        },
+      );
+    });
+
+    it('traces fetchLiveness API call with expected name', async () => {
+      const traceCallback = createTraceCallback();
+
+      await withController(
+        {
+          options: {
+            trace: traceCallback,
+          },
+        },
+        async ({ controller }) => {
+          nock(SENTINEL_API_BASE_URL_MAP[ethereumChainIdDec])
+            .get(`/network`)
+            .reply(200, createSuccessLivenessApiResponse());
+
+          await controller.fetchLiveness();
+
+          expect(traceCallback).toHaveBeenCalledWith(
+            { name: SmartTransactionsTraceName.FetchLiveness },
+            expect.any(Function),
+          );
+        },
+      );
+    });
+
+    it('returns correct result when tracing is enabled', async () => {
+      const traceCallback = createTraceCallback();
+
+      await withController(
+        {
+          options: {
+            trace: traceCallback,
+          },
+        },
+        async ({ controller }) => {
+          const apiUrl = API_BASE_URL;
+          const expectedResponse = createGetFeesApiResponse();
+          nock(apiUrl)
+            .post(`/networks/${ethereumChainIdDec}/getFees`)
+            .reply(200, expectedResponse);
+
+          const tradeTx = createUnsignedTransaction(ethereumChainIdDec);
+          const result = await controller.getFees(tradeTx);
+
+          expect(traceCallback).toHaveBeenCalled();
+          expect(result).toMatchObject({
+            tradeTxFees: expectedResponse.txs[0],
+            approvalTxFees: null,
+          });
+        },
+      );
+    });
+  });
 });
 
 type WithControllerCallback<ReturnValue> = ({

src/SmartTransactionsController.ts

@@ -9,6 +9,7 @@ import {
   safelyExecute,
   ChainId,
   isSafeDynamicKey,
+  type TraceCallback,
 } from '@metamask/controller-utils';
 import EthQuery from '@metamask/eth-query';
 import type {
@@ -27,7 +28,11 @@ import { TransactionStatus } from '@metamask/transaction-controller';
 import { BigNumber } from 'bignumber.js';
 import cloneDeep from 'lodash/cloneDeep';
 
-import { MetaMetricsEventCategory, MetaMetricsEventName } from './constants';
+import {
+  MetaMetricsEventCategory,
+  MetaMetricsEventName,
+  SmartTransactionsTraceName,
+} from './constants';
 import type {
   Fees,
   Hex,
@@ -206,6 +211,7 @@ type SmartTransactionsControllerOptions = {
   getMetaMetricsProps: () => Promise<MetaMetricsProps>;
   getFeatureFlags: () => FeatureFlags;
   updateTransaction: (transaction: TransactionMeta, note: string) => void;
+  trace?: TraceCallback;
 };
 
 export type SmartTransactionsControllerPollingInput = {
@@ -245,6 +251,8 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
 
   #updateTransaction: SmartTransactionsControllerOptions['updateTransaction'];
 
+  #trace: TraceCallback;
+
   /* istanbul ignore next */
   async #fetch(request: string, options?: RequestInit) {
     const fetchOptions = {
@@ -272,6 +280,7 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     getMetaMetricsProps,
     getFeatureFlags,
     updateTransaction,
+    trace,
   }: SmartTransactionsControllerOptions) {
     super({
       name: controllerName,
@@ -295,6 +304,7 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     this.#getMetaMetricsProps = getMetaMetricsProps;
     this.#getFeatureFlags = getFeatureFlags;
     this.#updateTransaction = updateTransaction;
+    this.#trace = trace ?? (((_request, fn) => fn?.()) as TraceCallback);
 
     this.initializeSmartTransactionsForChainId();
 
@@ -860,7 +870,7 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     const chainId = this.#getChainId({
       networkClientId: selectedNetworkClientId,
     });
-    const transactions = [];
+    const transactions: UnsignedTransaction[] = [];
     let unsignedTradeTransactionWithNonce;
     if (approvalTx) {
       const unsignedApprovalTransactionWithNonce =
@@ -880,14 +890,15 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
       );
     }
     transactions.push(unsignedTradeTransactionWithNonce);
-    const data = await this.#fetch(
-      getAPIRequestURL(APIType.GET_FEES, chainId),
-      {
-        method: 'POST',
-        body: JSON.stringify({
-          txs: transactions,
+    const data = await this.#trace(
+      { name: SmartTransactionsTraceName.GetFees },
+      async () =>
+        await this.#fetch(getAPIRequestURL(APIType.GET_FEES, chainId), {
+          method: 'POST',
+          body: JSON.stringify({
+            txs: transactions,
+          }),
         }),
-      },
     );
     let approvalTxFees: IndividualTxFees | null;
     let tradeTxFees: IndividualTxFees | null;
@@ -943,15 +954,19 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     const ethQuery = this.#getEthQuery({
       networkClientId: selectedNetworkClientId,
     });
-    const data = await this.#fetch(
-      getAPIRequestURL(APIType.SUBMIT_TRANSACTIONS, chainId),
-      {
-        method: 'POST',
-        body: JSON.stringify({
-          rawTxs: signedTransactions,
-          rawCancelTxs: signedCanceledTransactions,
-        }),
-      },
+    const data = await this.#trace(
+      { name: SmartTransactionsTraceName.SubmitTransactions },
+      async () =>
+        await this.#fetch(
+          getAPIRequestURL(APIType.SUBMIT_TRANSACTIONS, chainId),
+          {
+            method: 'POST',
+            body: JSON.stringify({
+              rawTxs: signedTransactions,
+              rawCancelTxs: signedCanceledTransactions,
+            }),
+          },
+        ),
     );
     const time = Date.now();
     let preTxBalance;
@@ -1089,10 +1104,14 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     } = {},
   ): Promise<void> {
     const chainId = this.#getChainId({ networkClientId });
-    await this.#fetch(getAPIRequestURL(APIType.CANCEL, chainId), {
-      method: 'POST',
-      body: JSON.stringify({ uuid }),
-    });
+    await this.#trace(
+      { name: SmartTransactionsTraceName.CancelTransaction },
+      async () =>
+        await this.#fetch(getAPIRequestURL(APIType.CANCEL, chainId), {
+          method: 'POST',
+          body: JSON.stringify({ uuid }),
+        }),
+    );
   }
 
   async fetchLiveness({
@@ -1103,8 +1122,10 @@ export default class SmartTransactionsController extends StaticIntervalPollingCo
     const chainId = this.#getChainId({ networkClientId });
     let liveness = false;
     try {
-      const response = await this.#fetch(
-        getAPIRequestURL(APIType.LIVENESS, chainId),
+      const response = await this.#trace(
+        { name: SmartTransactionsTraceName.FetchLiveness },
+        async () =>
+          await this.#fetch(getAPIRequestURL(APIType.LIVENESS, chainId)),
       );
       liveness = Boolean(response.smartTransactions);
     } catch (error) {

src/constants.ts

@@ -23,3 +23,10 @@ export enum MetaMetricsEventCategory {
   Transactions = 'Transactions',
   Navigation = 'Navigation',
 }
+
+export enum SmartTransactionsTraceName {
+  GetFees = 'Smart Transactions: Get Fees',
+  SubmitTransactions = 'Smart Transactions: Submit Transactions',
+  CancelTransaction = 'Smart Transactions: Cancel Transaction',
+  FetchLiveness = 'Smart Transactions: Fetch Liveness',
+}