refactor: MultichainAccountService, MultichainAccountWallet, MultichainAccountGroup performance and DevX improvements

packages/accounts-controller/src/AccountsController.ts

@@ -124,6 +124,11 @@ export type AccountsControllerGetAccountAction = {
   handler: AccountsController['getAccount'];
 };
 
+export type AccountsControllerGetAccountsAction = {
+  type: `${typeof controllerName}:getAccounts`;
+  handler: AccountsController['getAccounts'];
+};
+
 export type AccountsControllerUpdateAccountMetadataAction = {
   type: `${typeof controllerName}:updateAccountMetadata`;
   handler: AccountsController['updateAccountMetadata'];
@@ -145,6 +150,7 @@ export type AccountsControllerActions =
   | AccountsControllerGetSelectedAccountAction
   | AccountsControllerGetNextAvailableAccountNameAction
   | AccountsControllerGetAccountAction
+  | AccountsControllerGetAccountsAction
   | AccountsControllerGetSelectedMultichainAccountAction
   | AccountsControllerUpdateAccountMetadataAction;
 
@@ -303,6 +309,16 @@ export class AccountsController extends BaseController<
     return this.state.internalAccounts.accounts[accountId];
   }
 
+  /**
+   * Returns the internal account objects for the given account IDs, if they exist.
+   *
+   * @param accountIds - The IDs of the accounts to retrieve.
+   * @returns The internal account objects, or undefined if the account(s) do not exist.
+   */
+  getAccounts(accountIds: string[]): (InternalAccount | undefined)[] {
+    return accountIds.map((accountId) => this.getAccount(accountId));
+  }
+
   /**
    * Returns an array of all evm internal accounts.
    *

packages/accounts-controller/src/index.ts

@@ -13,6 +13,7 @@ export type {
   AccountsControllerGetAccountByAddressAction,
   AccountsControllerGetNextAvailableAccountNameAction,
   AccountsControllerGetAccountAction,
+  AccountsControllerGetAccountsAction,
   AccountsControllerUpdateAccountMetadataAction,
   AllowedActions,
   AccountsControllerActions,

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

@@ -199,7 +199,7 @@ describe('MultichainAccount', () => {
       expect(providers[1].createAccounts).not.toHaveBeenCalled();
     });
 
-    it('warns if provider alignment fails', async () => {
+    it('warns if alignment fails for a single provider', async () => {
       const groupIndex = 0;
       const { group, providers, wallet } = setup({
         groupIndex,
@@ -208,7 +208,7 @@ describe('MultichainAccount', () => {
 
       const consoleSpy = jest.spyOn(console, 'warn').mockImplementation();
       providers[1].createAccounts.mockRejectedValueOnce(
-        new Error('Unable to create accounts'),
+        new Error('Provider 2: Unable to create accounts'),
       );
 
       await group.alignAccounts();
@@ -219,7 +219,39 @@ describe('MultichainAccount', () => {
         groupIndex,
       });
       expect(consoleSpy).toHaveBeenCalledWith(
-        `Failed to fully align multichain account group for entropy ID: ${wallet.entropySource} and group index: ${groupIndex}, some accounts might be missing`,
+        `Failed to fully align multichain account group for entropy ID: ${wallet.entropySource} and group index: ${groupIndex}, some accounts might be missing. Provider threw the following error:\n- Error: Provider 2: Unable to create accounts`,
+      );
+    });
+
+    it('warns if alignment fails for multiple providers', async () => {
+      const groupIndex = 0;
+      const { group, providers, wallet } = setup({
+        groupIndex,
+        accounts: [[MOCK_WALLET_1_EVM_ACCOUNT], [], []],
+      });
+
+      const consoleSpy = jest.spyOn(console, 'warn').mockImplementation();
+      providers[1].createAccounts.mockRejectedValueOnce(
+        new Error('Provider 2: Unable to create accounts'),
+      );
+
+      providers[2].createAccounts.mockRejectedValueOnce(
+        new Error('Provider 3: Unable to create accounts'),
+      )
+
+      await group.align();
+
+      expect(providers[0].createAccounts).not.toHaveBeenCalled();
+      expect(providers[1].createAccounts).toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex,
+      });
+      expect(providers[2].createAccounts).toHaveBeenCalledWith({
+        entropySource: wallet.entropySource,
+        groupIndex,
+      });
+      expect(consoleSpy).toHaveBeenCalledWith(
+        `Failed to fully align multichain account group for entropy ID: ${wallet.entropySource} and group index: ${groupIndex}, some accounts might be missing. Providers threw the following errors:\n- Error: Provider 2: Unable to create accounts\n- Error: Provider 3: Unable to create accounts`,
       );
     });
   });

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

@@ -8,10 +8,14 @@ import type { Bip44Account } from '@metamask/account-api';
 import type { AccountSelector } from '@metamask/account-api';
 import { type KeyringAccount } from '@metamask/keyring-api';
 
+import type { ServiceState, StateKeys } from './MultichainAccountService';
 import type { MultichainAccountWallet } from './MultichainAccountWallet';
-import type { NamedAccountProvider } from './providers';
+import type { BaseBip44AccountProvider } from './providers';
 import type { MultichainAccountServiceMessenger } from './types';
 
+export type GroupState =
+  ServiceState[StateKeys['entropySource']][StateKeys['groupIndex']];
+
 /**
  * A multichain account group that holds multiple accounts.
  */
@@ -25,21 +29,14 @@ export class MultichainAccountGroup<
 
   readonly #groupIndex: number;
 
-  readonly #providers: NamedAccountProvider<Account>[];
+  readonly #providers: BaseBip44AccountProvider[];
 
-  readonly #providerToAccounts: Map<
-    NamedAccountProvider<Account>,
-    Account['id'][]
-  >;
+  readonly #providerToAccounts: Map<BaseBip44AccountProvider, Account['id'][]>;
 
-  readonly #accountToProvider: Map<
-    Account['id'],
-    NamedAccountProvider<Account>
-  >;
+  readonly #accountToProvider: Map<Account['id'], BaseBip44AccountProvider>;
 
   readonly #messenger: MultichainAccountServiceMessenger;
 
-  // eslint-disable-next-line @typescript-eslint/prefer-readonly
   #initialized = false;
 
   constructor({
@@ -50,7 +47,7 @@ export class MultichainAccountGroup<
   }: {
     groupIndex: number;
     wallet: MultichainAccountWallet<Account>;
-    providers: NamedAccountProvider<Account>[];
+    providers: BaseBip44AccountProvider[];
     messenger: MultichainAccountServiceMessenger;
   }) {
     this.#id = toMultichainAccountGroupId(wallet.id, groupIndex);
@@ -60,44 +57,37 @@ export class MultichainAccountGroup<
     this.#messenger = messenger;
     this.#providerToAccounts = new Map();
     this.#accountToProvider = new Map();
-
-    this.sync();
-    this.#initialized = true;
   }
 
   /**
-   * Force multichain account synchronization.
+   * Initialize the multichain account group and construct the internal representation of accounts.
+   *
+   * Note: This method can be called multiple times to update the group state.
    *
-   * This can be used if account providers got new accounts that the multichain
-   * account doesn't know about.
+   * @param groupState - The group state.
    */
-  sync(): void {
-    // Clear reverse mapping and re-construct it entirely based on the refreshed
-    // list of accounts from each providers.
-    this.#accountToProvider.clear();
-
+  init(groupState: GroupState) {
     for (const provider of this.#providers) {
-      // Filter account only for that index.
-      const accounts = [];
-      for (const account of provider.getAccounts()) {
-        if (
-          account.options.entropy.id === this.wallet.entropySource &&
-          account.options.entropy.groupIndex === this.groupIndex
-        ) {
-          // We only use IDs to always fetch the latest version of accounts.
-          accounts.push(account.id);
-        }
-      }
-      this.#providerToAccounts.set(provider, accounts);
+      const accountIds = groupState[provider.getName()];
 
-      // Reverse-mapping for fast indexing.
-      for (const id of accounts) {
-        this.#accountToProvider.set(id, provider);
+      if (accountIds) {
+        for (const accountId of accountIds) {
+          this.#accountToProvider.set(accountId, provider);
+        }
+        const providerAccounts = this.#providerToAccounts.get(provider);
+        if (!providerAccounts) {
+          this.#providerToAccounts.set(provider, accountIds);
+        } else {
+          providerAccounts.push(...accountIds);
+        }
+        // Add the accounts to the provider's internal list of account IDs
+        provider.addAccounts(accountIds);
       }
     }
 
-    // Emit update event when group is synced (only if initialized)
-    if (this.#initialized) {
+    if (!this.#initialized) {
+      this.#initialized = true;
+    } else {
       this.#messenger.publish(
         'MultichainAccountService:multichainAccountGroupUpdated',
         this,
@@ -167,14 +157,24 @@ export class MultichainAccountGroup<
           // If for some reason we cannot get this account from the provider, it
           // might means it has been deleted or something, so we just filter it
           // out.
-          allAccounts.push(account);
+          // We cast here because TS cannot infer the type of the account from the provider
+          allAccounts.push(account as Account);
         }
       }
     }
 
     return allAccounts;
   }
 
+  /**
+   * Gets the account IDs for this multichain account.
+   *
+   * @returns The account IDs.
+   */
+  getAccountIds(): Account['id'][] {
+    return [...this.#providerToAccounts.values()].flat();
+  }
+
   /**
    * Gets the account for a given account ID.
    *
@@ -189,7 +189,9 @@ export class MultichainAccountGroup<
     if (!provider) {
       return undefined;
     }
-    return provider.getAccount(id);
+
+    // We cast here because TS cannot infer the type of the account from the provider
+    return provider.getAccount(id) as Account;
   }
 
   /**
@@ -228,13 +230,33 @@ export class MultichainAccountGroup<
             groupIndex: this.groupIndex,
           });
         }
-        return Promise.resolve();
+        return Promise.reject(new Error('Already aligned'));
       }),
     );
 
+    const groupState = results.reduce<GroupState>((state, result, idx) => {
+      if (result.status === 'fulfilled') {
+        state[this.#providers[idx].getName()] = result.value.map(
+          (account) => account.id,
+        );
+      }
+      return state;
+    }, {});
+
+    // Update group state
+    this.init(groupState);
+
     if (results.some((result) => result.status === 'rejected')) {
+      const rejectedResults = results.filter(
+        (result) =>
+          result.status === 'rejected' && result.reason !== 'Already aligned',
+      ) as PromiseRejectedResult[];
+      const errors = rejectedResults
+        .map((result) => `- ${result.reason}`)
+        .join('\n');
+      const hasMultipleFailures = rejectedResults.length > 1;
       console.warn(
-        `Failed to fully align multichain account group for entropy ID: ${this.wallet.entropySource} and group index: ${this.groupIndex}, some accounts might be missing`,
+        `Failed to fully align multichain account group for entropy ID: ${this.wallet.entropySource} and group index: ${this.groupIndex}, some accounts might be missing. ${hasMultipleFailures ? 'Providers' : 'Provider'} threw the following ${hasMultipleFailures ? 'errors' : 'error'}:\n${errors}`,
       );
     }
   }