feat: add mutex + wallet.status support

Author: ccharly

packages/multichain-account-service/package.json

@@ -56,7 +56,8 @@
     "@metamask/snaps-sdk": "^9.0.0",
     "@metamask/snaps-utils": "^11.0.0",
     "@metamask/superstruct": "^3.1.0",
-    "@metamask/utils": "^11.4.2"
+    "@metamask/utils": "^11.4.2",
+    "async-mutex": "^0.5.0"
   },
   "devDependencies": {
     "@metamask/account-api": "^0.9.0",

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

@@ -13,8 +13,6 @@ import { MultichainAccountGroup } from './MultichainAccountGroup';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
 import type { MockAccountProvider } from './tests';
 import {
-  getMultichainAccountServiceMessenger,
-  getRootMessenger,
   MOCK_SNAP_ACCOUNT_2,
   MOCK_WALLET_1_BTC_P2TR_ACCOUNT,
   MOCK_WALLET_1_BTC_P2WPKH_ACCOUNT,

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

@@ -701,26 +701,50 @@ describe('MultichainAccountService', () => {
     });
 
     it('returns true during alignWallets and false after completion', async () => {
-      const { service } = setup({
+      const { service, messenger } = setup({
         accounts: [MOCK_HD_ACCOUNT_1],
       });
 
+      expect(service.getIsAlignmentInProgress()).toBe(false);
+
+      const mockWalletStatusChange = jest.fn().mockImplementationOnce(() => {
+        // When the alignment starts on any wallet, the service should report that
+        // properly.
+        expect(service.getIsAlignmentInProgress()).toBe(true);
+      });
+
+      messenger.subscribe(
+        'MultichainAccountService:walletStatusChange',
+        mockWalletStatusChange,
+      );
+
       const alignmentPromise = service.alignWallets();
-      expect(service.getIsAlignmentInProgress()).toBe(true);
 
       await alignmentPromise;
       expect(service.getIsAlignmentInProgress()).toBe(false);
     });
 
     it('returns true during alignWallet and false after completion', async () => {
-      const { service } = setup({
+      const { service, messenger } = setup({
         accounts: [MOCK_HD_ACCOUNT_1],
       });
 
+      expect(service.getIsAlignmentInProgress()).toBe(false);
+
+      const mockWalletStatusChange = jest.fn().mockImplementationOnce(() => {
+        // When the alignment starts on any wallet, the service should report that
+        // properly.
+        expect(service.getIsAlignmentInProgress()).toBe(true);
+      });
+
+      messenger.subscribe(
+        'MultichainAccountService:walletStatusChange',
+        mockWalletStatusChange,
+      );
+
       const alignmentPromise = service.alignWallet(
         MOCK_HD_KEYRING_1.metadata.id,
       );
-      expect(service.getIsAlignmentInProgress()).toBe(true);
 
       await alignmentPromise;
       expect(service.getIsAlignmentInProgress()).toBe(false);

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

@@ -151,7 +151,6 @@ export class MultichainAccountService {
         // (based on the accounts owned by each account providers).
         const wallet = new MultichainAccountWallet({
           entropySource,
-          messenger: this.#messenger,
           providers: this.#providers,
           messenger: this.#messenger,
         });
@@ -185,7 +184,6 @@ export class MultichainAccountService {
       // That's a new wallet.
       wallet = new MultichainAccountWallet({
         entropySource: account.options.entropy.id,
-        messenger: this.#messenger,
         providers: this.#providers,
         messenger: this.#messenger,
       });

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

@@ -18,8 +18,6 @@ import type { InternalAccount } from '@metamask/keyring-internal-api';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
 import type { MockAccountProvider } from './tests';
 import {
-  getMultichainAccountServiceMessenger,
-  getRootMessenger,
   MOCK_HD_ACCOUNT_1,
   MOCK_HD_KEYRING_1,
   MOCK_SNAP_ACCOUNT_2,
@@ -39,6 +37,7 @@ import type {
   AllowedEvents,
   MultichainAccountServiceActions,
   MultichainAccountServiceEvents,
+  MultichainAccountServiceMessenger,
 } from './types';
 
 function setup({
@@ -65,18 +64,21 @@ function setup({
 } = {}): {
   wallet: MultichainAccountWallet<Bip44Account<InternalAccount>>;
   providers: MockAccountProvider[];
+  messenger: MultichainAccountServiceMessenger;
 } {
   providers ??= accounts.map((providerAccounts) => {
     return setupAccountProvider({ accounts: providerAccounts });
   });
 
+  const serviceMessenger = getMultichainAccountServiceMessenger(messenger);
+
   const wallet = new MultichainAccountWallet<Bip44Account<InternalAccount>>({
     entropySource,
     providers,
-    messenger: getMultichainAccountServiceMessenger(messenger),
+    messenger: serviceMessenger,
   });
 
-  return { wallet, providers };
+  return { wallet, providers, messenger: serviceMessenger };
 }
 
 describe('MultichainAccountWallet', () => {
@@ -440,85 +442,28 @@ describe('MultichainAccountWallet', () => {
     });
 
     it('returns true during alignment and false after completion', async () => {
-      const { wallet } = setup();
-
-      // Start alignment (don't await yet)
-      const alignmentPromise = wallet.alignGroups();
+      const { wallet, messenger } = setup();
 
-      // Check if alignment is in progress
-      expect(wallet.getIsAlignmentInProgress()).toBe(true);
-
-      // Wait for completion
-      await alignmentPromise;
-
-      // Should be false after completion
       expect(wallet.getIsAlignmentInProgress()).toBe(false);
-    });
-  });
-
-  describe('concurrent alignment prevention', () => {
-    it('prevents concurrent alignGroups calls', async () => {
-      // Setup with EVM account in group 0, Sol account in group 1 (missing group 0)
-      const mockEvmAccount = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
-        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
-        .withGroupIndex(0)
-        .get();
-      const mockSolAccount = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
-        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
-        .withGroupIndex(1)
-        .get();
-      const { wallet, providers } = setup({
-        accounts: [[mockEvmAccount], [mockSolAccount]],
-      });
-
-      // Make provider createAccounts slow to ensure concurrency
-      providers[1].createAccounts.mockImplementation(
-        () => new Promise((resolve) => setTimeout(() => resolve([]), 50)),
-      );
-
-      // Start first alignment
-      const firstAlignment = wallet.alignGroups();
 
-      // Start second alignment while first is still running
-      const secondAlignment = wallet.alignGroups();
-
-      // Both should complete without error
-      await Promise.all([firstAlignment, secondAlignment]);
-
-      // Provider should only be called once (not twice due to concurrency protection)
-      expect(providers[1].createAccounts).toHaveBeenCalledTimes(1);
-    });
-
-    it('prevents concurrent alignGroup calls', async () => {
-      // Setup with EVM account in group 0, Sol account in group 1 (missing group 0)
-      const mockEvmAccount = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
-        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
-        .withGroupIndex(0)
-        .get();
-      const mockSolAccount = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
-        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
-        .withGroupIndex(1)
-        .get();
-      const { wallet, providers } = setup({
-        accounts: [[mockEvmAccount], [mockSolAccount]],
+      const mockWalletStatusChange = jest.fn().mockImplementationOnce(() => {
+        // Listen for wallet status change and check that the alignment is in progress.
+        expect(wallet.getIsAlignmentInProgress()).toBe(true);
       });
 
-      // Make provider createAccounts slow to ensure concurrency
-      providers[1].createAccounts.mockImplementation(
-        () => new Promise((resolve) => setTimeout(() => resolve([]), 50)),
+      messenger.subscribe(
+        'MultichainAccountService:walletStatusChange',
+        mockWalletStatusChange,
       );
 
-      // Start first alignment
-      const firstAlignment = wallet.alignGroup(0);
-
-      // Start second alignment while first is still running
-      const secondAlignment = wallet.alignGroup(0);
+      // Start alignment (don't await yet)
+      const alignmentPromise = wallet.alignGroups();
 
-      // Both should complete without error
-      await Promise.all([firstAlignment, secondAlignment]);
+      // Wait for completion
+      await alignmentPromise;
 
-      // Provider should only be called once (not twice due to concurrency protection)
-      expect(providers[1].createAccounts).toHaveBeenCalledTimes(1);
+      // Should be false after completion
+      expect(wallet.getIsAlignmentInProgress()).toBe(false);
     });
   });
 
@@ -551,6 +496,7 @@ describe('MultichainAccountWallet', () => {
       jest.useFakeTimers();
       const discovery = wallet.discoverAndCreateAccounts();
       // Allow fast provider microtasks to run and advance maxGroupIndex first
+      await Promise.resolve(); // Mutex lock.
       await Promise.resolve();
       await Promise.resolve();
       jest.advanceTimersByTime(100);