feat: add tracing to multichain account service

packages/multichain-account-service/CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Add optional tracing configuration ([#7006](https://github.com/MetaMask/core/pull/7006))
+  - For now, only the account discovery is being traced.
 - Limit Bitcoin and Tron providers to 3 concurrent account creations by default when creating multichain account groups ([#7052](https://github.com/MetaMask/core/pull/7052))
 
 ## [2.1.0]

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

@@ -220,10 +220,12 @@ describe('MultichainAccountService', () => {
       expect(mocks.EvmAccountProvider.constructor).toHaveBeenCalledWith(
         messenger,
         providerConfigs?.[EvmAccountProvider.NAME],
+        expect.any(Function), // TraceCallback
       );
       expect(mocks.SolAccountProvider.constructor).toHaveBeenCalledWith(
         messenger,
         providerConfigs?.[SolAccountProvider.NAME],
+        expect.any(Function), // TraceCallback
       );
     });
 
@@ -254,10 +256,12 @@ describe('MultichainAccountService', () => {
       expect(mocks.EvmAccountProvider.constructor).toHaveBeenCalledWith(
         messenger,
         undefined,
+        expect.any(Function), // TraceCallback
       );
       expect(mocks.SolAccountProvider.constructor).toHaveBeenCalledWith(
         messenger,
         providerConfigs?.[SolAccountProvider.NAME],
+        expect.any(Function), // TraceCallback
       );
     });
   });
@@ -563,10 +567,20 @@ describe('MultichainAccountService', () => {
       );
 
       // Should emit updated event for the existing group
+      expect(publishSpy).toHaveBeenCalled();
       expect(publishSpy).toHaveBeenCalledWith(
         'MultichainAccountService:multichainAccountGroupUpdated',
-        expect.any(Object),
+        expect.objectContaining({
+          groupIndex: 0,
+        }),
       );
+
+      const emittedGroup = publishSpy.mock.calls[0][1];
+      expect(emittedGroup).toBeDefined();
+      expect(emittedGroup).toHaveProperty('groupIndex', 0);
+      expect(emittedGroup).toHaveProperty('getAccounts');
+      expect(emittedGroup).toHaveProperty('select');
+      expect(emittedGroup).toHaveProperty('get');
     });
 
     it('creates new detected wallets and update reverse mapping on AccountsController:accountAdded', () => {

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

@@ -12,6 +12,7 @@ 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';
@@ -28,7 +29,10 @@ import {
   SolAccountProvider,
   type SolAccountProviderConfig,
 } from './providers/SolAccountProvider';
-import type { MultichainAccountServiceMessenger } from './types';
+import type {
+  MultichainAccountServiceConfig,
+  MultichainAccountServiceMessenger,
+} from './types';
 
 export const serviceName = 'MultichainAccountService';
 
@@ -42,6 +46,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. */
@@ -81,28 +86,35 @@ 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();
 
+    // Pass trace callback directly to preserve original 'this' context
+    // This avoids binding the callback to the MultichainAccountService instance
+    const traceCallback = config?.trace ?? traceFallback;
+
     // TODO: Rely on keyring capabilities once the keyring API is used by all keyrings.
     this.#providers = [
       new EvmAccountProvider(
         this.#messenger,
         providerConfigs?.[EvmAccountProvider.NAME],
+        traceCallback,
       ),
       new AccountProviderWrapper(
         this.#messenger,
         new SolAccountProvider(
           this.#messenger,
           providerConfigs?.[SolAccountProvider.NAME],
+          traceCallback,
         ),
       ),
       // Custom account providers that can be provided by the MetaMask client.

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 { traceFallback } from './traces';
+import { TraceName } from '../constants/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/BtcAccountProvider.test.ts

@@ -8,7 +8,11 @@ import type {
 } from '@metamask/keyring-internal-api';
 
 import { AccountProviderWrapper } from './AccountProviderWrapper';
-import { BtcAccountProvider } from './BtcAccountProvider';
+import {
+  BTC_ACCOUNT_PROVIDER_NAME,
+  BtcAccountProvider,
+} from './BtcAccountProvider';
+import { TraceName } from '../constants/traces';
 import {
   getMultichainAccountServiceMessenger,
   getRootMessenger,
@@ -324,4 +328,129 @@ describe('BtcAccountProvider', () => {
 
     expect(discovered).toStrictEqual([]);
   });
+
+  describe('trace functionality', () => {
+    it('calls trace callback during account discovery', async () => {
+      const mockTrace = jest.fn().mockImplementation(async (request, fn) => {
+        expect(request.name).toBe(TraceName.SnapDiscoverAccounts);
+        expect(request.data).toStrictEqual({
+          provider: BTC_ACCOUNT_PROVIDER_NAME,
+        });
+        return await fn();
+      });
+
+      const { messenger, mocks } = setup({
+        accounts: [],
+      });
+
+      // Simulate one discovered account at the requested index.
+      mocks.handleRequest.mockReturnValue([MOCK_BTC_P2TR_DISCOVERED_ACCOUNT_1]);
+
+      const multichainMessenger =
+        getMultichainAccountServiceMessenger(messenger);
+      const btcProvider = new BtcAccountProvider(
+        multichainMessenger,
+        undefined,
+        mockTrace,
+      );
+      const provider = new AccountProviderWrapper(
+        multichainMessenger,
+        btcProvider,
+      );
+
+      const discovered = await provider.discoverAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(discovered).toHaveLength(1);
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+    });
+
+    it('uses fallback trace when no trace callback is provided', async () => {
+      const { provider, mocks } = setup({
+        accounts: [],
+      });
+
+      mocks.handleRequest.mockReturnValue([MOCK_BTC_P2TR_DISCOVERED_ACCOUNT_1]);
+
+      const discovered = await provider.discoverAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(discovered).toHaveLength(1);
+      // No trace errors, fallback trace should be used silently
+    });
+
+    it('trace callback is called even when discovery returns empty results', async () => {
+      const mockTrace = jest.fn().mockImplementation(async (request, fn) => {
+        expect(request.name).toBe(TraceName.SnapDiscoverAccounts);
+        expect(request.data).toStrictEqual({
+          provider: BTC_ACCOUNT_PROVIDER_NAME,
+        });
+        return await fn();
+      });
+
+      const { messenger, mocks } = setup({
+        accounts: [],
+      });
+
+      mocks.handleRequest.mockReturnValue([]);
+
+      const multichainMessenger =
+        getMultichainAccountServiceMessenger(messenger);
+      const btcProvider = new BtcAccountProvider(
+        multichainMessenger,
+        undefined,
+        mockTrace,
+      );
+      const provider = new AccountProviderWrapper(
+        multichainMessenger,
+        btcProvider,
+      );
+
+      const discovered = await provider.discoverAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(discovered).toStrictEqual([]);
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+    });
+
+    it('trace callback receives error when discovery fails', async () => {
+      const mockError = new Error('Discovery failed');
+      const mockTrace = jest.fn().mockImplementation(async (_request, fn) => {
+        return await fn();
+      });
+
+      const { messenger, mocks } = setup({
+        accounts: [],
+      });
+
+      mocks.handleRequest.mockRejectedValue(mockError);
+
+      const multichainMessenger =
+        getMultichainAccountServiceMessenger(messenger);
+      const btcProvider = new BtcAccountProvider(
+        multichainMessenger,
+        undefined,
+        mockTrace,
+      );
+      const provider = new AccountProviderWrapper(
+        multichainMessenger,
+        btcProvider,
+      );
+
+      await expect(
+        provider.discoverAccounts({
+          entropySource: MOCK_HD_KEYRING_1.metadata.id,
+          groupIndex: 0,
+        }),
+      ).rejects.toThrow(mockError);
+
+      expect(mockTrace).toHaveBeenCalledTimes(1);
+    });
+  });
 });