feat: add alignment methods to group, wallet and service classes

packages/multichain-account-service/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Allow for multichain account group alignment through the `align` method ([#6326](https://github.com/MetaMask/core/pull/6326))
+  - You can now call alignment from the group, wallet and service levels.
+
 ### Changed
 
 - Bump `@metamask/base-controller` from `^8.0.1` to `^8.1.0` ([#6284](https://github.com/MetaMask/core/pull/6284))

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

@@ -146,4 +146,38 @@ describe('MultichainAccount', () => {
       expect(group.select({ scopes: [SolScope.Mainnet] })).toStrictEqual([]);
     });
   });
+
+  describe('align', () => {
+    it('creates missing accounts only for providers with no accounts', async () => {
+      const groupIndex = 0;
+      const { group, providers, wallet } = setup({
+        groupIndex,
+        accounts: [
+          [MOCK_WALLET_1_EVM_ACCOUNT], // provider[0] already has group 0
+          [], // provider[1] missing group 0
+        ],
+      });
+
+      await group.align();
+
+      expect(providers[0].createAccounts).not.toHaveBeenCalled();
+      expect(providers[1].createAccounts).toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex,
+      });
+    });
+
+    it('does nothing when already aligned', async () => {
+      const groupIndex = 0;
+      const { group, providers } = setup({
+        groupIndex,
+        accounts: [[MOCK_WALLET_1_EVM_ACCOUNT], [MOCK_WALLET_1_SOL_ACCOUNT]],
+      });
+
+      await group.align();
+
+      expect(providers[0].createAccounts).not.toHaveBeenCalled();
+      expect(providers[1].createAccounts).not.toHaveBeenCalled();
+    });
+  });
 });

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

@@ -188,4 +188,24 @@ export class MultichainAccountGroup<
   select(selector: AccountSelector<Account>): Account[] {
     return select(this.getAccounts(), selector);
   }
+
+  /**
+   * Align the multichain account group.
+   *
+   * This will create accounts for providers that don't have any accounts yet.
+   */
+  async align(): Promise<void> {
+    await Promise.all(
+      this.#providers.map((provider) => {
+        const accounts = this.#providerToAccounts.get(provider);
+        if (!accounts || accounts.length === 0) {
+          return provider.createAccounts({
+            entropySource: this.wallet.entropySource,
+            groupIndex: this.groupIndex,
+          });
+        }
+        return Promise.resolve();
+      }),
+    );
+  }
 }

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

@@ -15,6 +15,7 @@ import {
   MOCK_HD_ACCOUNT_2,
   MOCK_SNAP_ACCOUNT_1,
   MOCK_SNAP_ACCOUNT_2,
+  MOCK_SOL_ACCOUNT_1,
   MockAccountBuilder,
 } from './tests';
 import {
@@ -574,6 +575,57 @@ describe('MultichainAccountService', () => {
     });
   });
 
+  describe('alignWallets', () => {
+    it('aligns all multichain account wallets', async () => {
+      const mockEvmAccount1 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const mockSolAccount1 = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_2.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const { service, mocks } = setup({
+        accounts: [mockEvmAccount1, mockSolAccount1],
+      });
+
+      await service.alignWallets();
+
+      expect(mocks.EvmAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_2.metadata.id,
+        groupIndex: 0,
+      });
+      expect(mocks.SolAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+    });
+  });
+
+  describe('alignWallet', () => {
+    it('aligns a specific multichain account wallet', async () => {
+      const mockEvmAccount1 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const mockSolAccount1 = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_2.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const { service, mocks } = setup({
+        accounts: [mockEvmAccount1, mockSolAccount1],
+      });
+
+      await service.alignWallet(MOCK_HD_KEYRING_1.metadata.id);
+
+      expect(mocks.SolAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+      expect(mocks.EvmAccountProvider.createAccounts).not.toHaveBeenCalled();
+    });
+  });
+
   describe('actions', () => {
     it('gets a multichain account with MultichainAccountService:getMultichainAccount', () => {
       const accounts = [MOCK_HD_ACCOUNT_1];
@@ -647,5 +699,55 @@ describe('MultichainAccountService', () => {
       expect(firstGroup.getAccounts()).toHaveLength(1);
       expect(firstGroup.getAccounts()[0]).toStrictEqual(MOCK_HD_ACCOUNT_1);
     });
+
+    it('aligns a multichain account wallet with MultichainAccountService:alignWallet', async () => {
+      const mockEvmAccount1 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const mockSolAccount1 = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_2.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const { messenger, mocks } = setup({
+        accounts: [mockEvmAccount1, mockSolAccount1],
+      });
+
+      await messenger.call(
+        'MultichainAccountService:alignWallet',
+        MOCK_HD_KEYRING_1.metadata.id,
+      );
+
+      expect(mocks.SolAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+      expect(mocks.EvmAccountProvider.createAccounts).not.toHaveBeenCalled();
+    });
+
+    it('aligns all multichain account wallets with MultichainAccountService:alignWallets', async () => {
+      const mockEvmAccount1 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const mockSolAccount1 = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_2.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const { messenger, mocks } = setup({
+        accounts: [mockEvmAccount1, mockSolAccount1],
+      });
+
+      await messenger.call('MultichainAccountService:alignWallets');
+
+      expect(mocks.EvmAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_2.metadata.id,
+        groupIndex: 0,
+      });
+      expect(mocks.SolAccountProvider.createAccounts).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+    });
   });
 });

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

@@ -105,6 +105,14 @@ export class MultichainAccountService {
       'MultichainAccountService:createMultichainAccountGroup',
       (...args) => this.createMultichainAccountGroup(...args),
     );
+    this.#messenger.registerActionHandler(
+      'MultichainAccountService:alignWallets',
+      (...args) => this.alignWallets(...args),
+    );
+    this.#messenger.registerActionHandler(
+      'MultichainAccountService:alignWallet',
+      (...args) => this.alignWallet(...args),
+    );
   }
 
   /**
@@ -350,4 +358,22 @@ export class MultichainAccountService {
       groupIndex,
     );
   }
+
+  /**
+   * Align all multichain account wallets.
+   */
+  async alignWallets(): Promise<void> {
+    const wallets = this.getMultichainAccountWallets();
+    await Promise.all(wallets.map((w) => w.alignGroups()));
+  }
+
+  /**
+   * Align a specific multichain account wallet.
+   *
+   * @param entropySource - The entropy source of the multichain account wallet.
+   */
+  async alignWallet(entropySource: EntropySourceId): Promise<void> {
+    const wallet = this.getMultichainAccountWallet({ entropySource });
+    await wallet.alignGroups();
+  }
 }

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

@@ -346,4 +346,67 @@ describe('MultichainAccountWallet', () => {
       expect(internalAccounts[1].type).toBe(SolAccountType.DataAccount);
     });
   });
+
+  describe('alignGroups', () => {
+    it('creates missing accounts only for providers with no accounts associated with a particular group index', async () => {
+      const mockEvmAccount1 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const mockEvmAccount2 = MockAccountBuilder.from(MOCK_HD_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(1)
+        .get();
+      const mockSolAccount = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1)
+        .withEntropySource(MOCK_HD_KEYRING_1.metadata.id)
+        .withGroupIndex(0)
+        .get();
+      const { wallet, providers } = setup({
+        accounts: [[mockEvmAccount1, mockEvmAccount2], [mockSolAccount]],
+      });
+
+      await wallet.alignGroups();
+
+      // EVM provider already has group 0 and 1; should not be called.
+      expect(providers[0].createAccounts).not.toHaveBeenCalled();
+
+      // Sol provider is missing group 1; should be called to create it.
+      expect(providers[1].createAccounts).toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex: 1,
+      });
+    });
+  });
+
+  describe('alignGroup', () => {
+    it('aligns a specific multichain account group', async () => {
+      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]],
+      });
+
+      await wallet.alignGroup(0);
+
+      // EVM provider already has group 0; should not be called.
+      expect(providers[0].createAccounts).not.toHaveBeenCalled();
+
+      // Sol provider is missing group 0; should be called to create it.
+      expect(providers[1].createAccounts).toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex: 0,
+      });
+
+      expect(providers[1].createAccounts).not.toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex: 1,
+      });
+    });
+  });
 });

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

@@ -306,4 +306,24 @@ export class MultichainAccountWallet<
   > {
     return this.createMultichainAccountGroup(this.getNextGroupIndex());
   }
+
+  /**
+   * Align all multichain account groups.
+   */
+  async alignGroups(): Promise<void> {
+    const groups = this.getMultichainAccountGroups();
+    await Promise.all(groups.map((g) => g.align()));
+  }
+
+  /**
+   * Align a specific multichain account group.
+   *
+   * @param groupIndex - The group index to align.
+   */
+  async alignGroup(groupIndex: number): Promise<void> {
+    const group = this.getMultichainAccountGroup(groupIndex);
+    if (group) {
+      await group.align();
+    }
+  }
 }

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

@@ -9,7 +9,7 @@ import type {
 import type { Hex } from '@metamask/utils';
 
 import {
-  assertIsBip44Account,
+  assertAreBip44Accounts,
   BaseAccountProvider,
 } from './BaseAccountProvider';
 
@@ -69,9 +69,11 @@ export class EvmAccountProvider extends BaseAccountProvider {
 
     // We MUST have the associated internal account.
     assertInternalAccountExists(account);
-    assertIsBip44Account(account);
 
-    return [account];
+    const accountsArray = [account];
+    assertAreBip44Accounts(accountsArray);
+
+    return accountsArray;
   }
 
   async discoverAndCreateAccounts(_: {