feat: add tracing to multichain account service

packages/multichain-account-service/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Adds an optional trace argument to MultichainAccountService and MultichainAccountWallet ([#7006](https://github.com/MetaMask/core/pull/7006))
+
 ## [2.0.0]
 
 ### Changed

packages/multichain-account-service/src/MultichainAccountService.test.ts

@@ -226,6 +226,30 @@ describe('MultichainAccountService', () => {
       );
     });
 
+    it('accepts trace callback in config', () => {
+      const mockTrace = jest.fn().mockImplementation((_request, fn) => fn?.());
+      const config = { trace: mockTrace };
+
+      // Test that the service can be constructed with a trace config
+      const rootMessenger = getRootMessenger();
+      const messenger = getMultichainAccountServiceMessenger(rootMessenger);
+
+      // Mock the required handlers for initialization
+      rootMessenger.registerActionHandler(
+        'KeyringController:getState',
+        jest.fn().mockReturnValue({ isUnlocked: true, keyrings: [] }),
+      );
+
+      const serviceWithTrace = new MultichainAccountService({
+        messenger,
+        config,
+      });
+
+      serviceWithTrace.init();
+
+      expect(serviceWithTrace).toBeDefined();
+    });
+
     it('allows optional configs for some providers', () => {
       const providerConfigs: MultichainAccountServiceOptions['providerConfigs'] =
         {
@@ -1090,6 +1114,57 @@ describe('MultichainAccountService', () => {
       (solProvider.isAccountCompatible as jest.Mock).mockReturnValue(false);
       expect(wrapper.isAccountCompatible(MOCK_HD_ACCOUNT_1)).toBe(false);
     });
+
+    it('calls trace callback when discoverAccounts() is enabled', async () => {
+      const mockTrace = jest.fn().mockImplementation(async (request, fn) => {
+        expect(request.name).toBe('Snap Discover Accounts');
+        expect(request.data).toStrictEqual({
+          providerName: solProvider.getName(),
+        });
+        return await fn();
+      });
+
+      const options = {
+        entropySource: MOCK_HD_ACCOUNT_1.options.entropy.id,
+        groupIndex: 0,
+      };
+
+      // Create wrapper with custom trace callback
+      const wrapperWithTrace = new AccountProviderWrapper(
+        getMultichainAccountServiceMessenger(
+          setup({ accounts: [MOCK_HD_ACCOUNT_1] }).rootMessenger,
+        ),
+        solProvider,
+        mockTrace,
+      );
+
+      (solProvider.discoverAccounts as jest.Mock).mockResolvedValue([
+        MOCK_HD_ACCOUNT_1,
+      ]);
+
+      const result = await wrapperWithTrace.discoverAccounts(options);
+
+      expect(result).toStrictEqual([MOCK_HD_ACCOUNT_1]);
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+      expect(solProvider.discoverAccounts).toHaveBeenCalledWith(options);
+    });
+
+    it('uses fallback trace when no trace callback is provided', async () => {
+      const options = {
+        entropySource: MOCK_HD_ACCOUNT_1.options.entropy.id,
+        groupIndex: 0,
+      };
+
+      (solProvider.discoverAccounts as jest.Mock).mockResolvedValue([
+        MOCK_HD_ACCOUNT_1,
+      ]);
+
+      // Wrapper without trace callback should use fallback
+      const result = await wrapper.discoverAccounts(options);
+
+      expect(result).toStrictEqual([MOCK_HD_ACCOUNT_1]);
+      expect(solProvider.discoverAccounts).toHaveBeenCalledWith(options);
+    });
   });
 
   describe('createMultichainAccountWallet', () => {

packages/multichain-account-service/src/MultichainAccountService.ts

@@ -6,12 +6,14 @@ import type {
   MultichainAccountWalletId,
   Bip44Account,
 } from '@metamask/account-api';
+import type { TraceCallback } from '@metamask/controller-utils';
 import type { HdKeyring } from '@metamask/eth-hd-keyring';
 import { mnemonicPhraseToBytes } from '@metamask/key-tree';
 import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api';
 import { KeyringTypes } from '@metamask/keyring-controller';
 import { areUint8ArraysEqual } from '@metamask/utils';
 
+import { traceFallback } from './analytics';
 import { projectLogger as log } from './logger';
 import type { MultichainAccountGroup } from './MultichainAccountGroup';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
@@ -26,7 +28,10 @@ import {
 } from './providers/AccountProviderWrapper';
 import { EvmAccountProvider } from './providers/EvmAccountProvider';
 import { SolAccountProvider } from './providers/SolAccountProvider';
-import type { MultichainAccountServiceMessenger } from './types';
+import type {
+  MultichainAccountServiceConfig,
+  MultichainAccountServiceMessenger,
+} from './types';
 
 export const serviceName = 'MultichainAccountService';
 
@@ -40,6 +45,7 @@ export type MultichainAccountServiceOptions = {
     [EvmAccountProvider.NAME]?: EvmAccountProviderConfig;
     [SolAccountProvider.NAME]?: SolAccountProviderConfig;
   };
+  config?: MultichainAccountServiceConfig;
 };
 
 /** Reverse mapping object used to map account IDs and their wallet/multichain account. */
@@ -66,6 +72,8 @@ export class MultichainAccountService {
     AccountContext<Bip44Account<KeyringAccount>>
   >;
 
+  readonly #trace: TraceCallback;
+
   /**
    * The name of the service.
    */
@@ -79,16 +87,18 @@ export class MultichainAccountService {
    * MultichainAccountService.
    * @param options.providers - Optional list of account
    * @param options.providerConfigs - Optional provider configs
-   * providers.
+   * @param options.config - Optional config.
    */
   constructor({
     messenger,
     providers = [],
     providerConfigs,
+    config,
   }: MultichainAccountServiceOptions) {
     this.#messenger = messenger;
     this.#wallets = new Map();
     this.#accountIdToContext = new Map();
+    this.#trace = config?.trace ?? traceFallback;
 
     // TODO: Rely on keyring capabilities once the keyring API is used by all keyrings.
     this.#providers = [
@@ -102,6 +112,7 @@ export class MultichainAccountService {
           this.#messenger,
           providerConfigs?.[SolAccountProvider.NAME],
         ),
+        this.#trace,
       ),
       // Custom account providers that can be provided by the MetaMask client.
       ...providers,

packages/multichain-account-service/src/analytics/index.ts

@@ -0,0 +1 @@
+export * from './traces';

packages/multichain-account-service/src/analytics/traces.test.ts

@@ -0,0 +1,76 @@
+import type { TraceRequest } from '@metamask/controller-utils';
+import { TraceName } from '../constants/traces';
+
+import { traceFallback } from './traces';
+
+describe('MultichainAccountService - Traces', () => {
+  describe('TraceName', () => {
+    it('contains expected trace names', () => {
+      expect(TraceName).toStrictEqual({
+        SnapDiscoverAccounts: 'Snap Discover Accounts',
+        EvmDiscoverAccounts: 'EVM Discover Accounts',
+      });
+    });
+  });
+
+  describe('traceFallback', () => {
+    let mockTraceRequest: TraceRequest;
+
+    beforeEach(() => {
+      mockTraceRequest = {
+        name: TraceName.SnapDiscoverAccounts,
+        id: 'trace-id-123',
+        tags: {},
+      };
+    });
+
+    it('returns undefined when no function is provided', async () => {
+      const result = await traceFallback(mockTraceRequest);
+
+      expect(result).toBeUndefined();
+    });
+
+    it('executes the provided function and return its result', async () => {
+      const mockResult = 'test-result';
+      const mockFn = jest.fn().mockReturnValue(mockResult);
+
+      const result = await traceFallback(mockTraceRequest, mockFn);
+
+      expect(mockFn).toHaveBeenCalledTimes(1);
+      expect(mockFn).toHaveBeenCalledWith();
+      expect(result).toBe(mockResult);
+    });
+
+    it('executes async function and return its result', async () => {
+      const mockResult = { data: 'async-result' };
+      const mockAsyncFn = jest.fn().mockResolvedValue(mockResult);
+
+      const result = await traceFallback(mockTraceRequest, mockAsyncFn);
+
+      expect(mockAsyncFn).toHaveBeenCalledTimes(1);
+      expect(result).toBe(mockResult);
+    });
+
+    it('handles function that throws an error', async () => {
+      const mockError = new Error('Test error');
+      const mockFn = jest.fn().mockImplementation(() => {
+        throw mockError;
+      });
+
+      await expect(traceFallback(mockTraceRequest, mockFn)).rejects.toThrow(
+        mockError,
+      );
+      expect(mockFn).toHaveBeenCalledTimes(1);
+    });
+
+    it('handles function that returns a rejected promise', async () => {
+      const mockError = new Error('Async error');
+      const mockFn = jest.fn().mockRejectedValue(mockError);
+
+      await expect(traceFallback(mockTraceRequest, mockFn)).rejects.toThrow(
+        mockError,
+      );
+      expect(mockFn).toHaveBeenCalledTimes(1);
+    });
+  });
+});

packages/multichain-account-service/src/analytics/traces.ts

@@ -0,0 +1,25 @@
+import type {
+  TraceCallback,
+  TraceContext,
+  TraceRequest,
+} from '@metamask/controller-utils';
+
+/**
+ * Fallback function for tracing.
+ * This function is used when no specific trace function is provided.
+ * It executes the provided function in a trace context if available.
+ *
+ * @param _request - The trace request containing additional data and context.
+ * @param fn - The function to execute within the trace context.
+ * @returns A promise that resolves to the result of the executed function.
+ * If no function is provided, it resolves to undefined.
+ */
+export const traceFallback: TraceCallback = async <ReturnType>(
+  _request: TraceRequest,
+  fn?: (context?: TraceContext) => ReturnType,
+): Promise<ReturnType> => {
+  if (!fn) {
+    return undefined as ReturnType;
+  }
+  return await Promise.resolve(fn());
+};

packages/multichain-account-service/src/constants/traces.ts

@@ -0,0 +1,4 @@
+export enum TraceName {
+  'SnapDiscoverAccounts' = 'Snap Discover Accounts',
+  'EvmDiscoverAccounts' = 'EVM Discover Accounts',
+}

packages/multichain-account-service/src/providers/AccountProviderWrapper.ts

@@ -1,7 +1,10 @@
 import type { Bip44Account } from '@metamask/account-api';
+import type { TraceCallback } from '@metamask/controller-utils';
 import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api';
 
 import { BaseBip44AccountProvider } from './BaseBip44AccountProvider';
+import { traceFallback } from '../analytics';
+import { TraceName } from '../constants/traces';
 import type { MultichainAccountServiceMessenger } from '../types';
 
 /**
@@ -13,12 +16,16 @@ export class AccountProviderWrapper extends BaseBip44AccountProvider {
 
   private readonly provider: BaseBip44AccountProvider;
 
+  readonly #trace: TraceCallback;
+
   constructor(
     messenger: MultichainAccountServiceMessenger,
     provider: BaseBip44AccountProvider,
+    trace?: TraceCallback,
   ) {
     super(messenger);
     this.provider = provider;
+    this.#trace = trace ?? traceFallback;
   }
 
   override getName(): string {
@@ -106,7 +113,15 @@ export class AccountProviderWrapper extends BaseBip44AccountProvider {
     if (!this.isEnabled) {
       return [];
     }
-    return this.provider.discoverAccounts(options);
+    return this.#trace(
+      {
+        name: TraceName.SnapDiscoverAccounts,
+        data: {
+          providerName: this.getName(),
+        },
+      },
+      async () => this.provider.discoverAccounts(options),
+    );
   }
 }