refactor(multichain-account-service): Improved performance across package classes and improved error messages

packages/accounts-controller/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add a `getAccounts` method (and its associated action) that is the plural version of `getAccount` ([#6654](https://github.com/MetaMask/core/pull/6708))
+  - This method is added to primarily be consumed in the `MultichainAccountService`.
+
 ### Changed
 
 - Bump `@metamask/utils` from `^11.4.2` to `^11.8.1` ([#6588](https://github.com/MetaMask/core/pull/6588), [#6708](https://github.com/MetaMask/core/pull/6708))

packages/accounts-controller/src/AccountsController.test.ts

@@ -20,7 +20,7 @@
 import { SnapStatus } from '@metamask/snaps-utils';
 import type { CaipChainId } from '@metamask/utils';
 import type { V4Options } from 'uuid';
 import * as uuid from 'uuid';
 
 import type {
   AccountsControllerActions,
@@ -2777,6 +2777,26 @@
     });
   });
 
+  describe('getAccounts', () => {
+    it('returns a list of accounts based on the given account IDs', () => {
+      const { accountsController } = setupAccountsController({
+        initialState: {
+          internalAccounts: {
+            accounts: {
+              [mockAccount.id]: mockAccount,
+              [mockAccount2.id]: mockAccount2,
+            },
+            selectedAccount: mockAccount.id,
+          },
+        },
+      });
+
+      const result = accountsController.getAccounts([mockAccount.id]);
+
+      expect(result).toStrictEqual([mockAccount]);
+    });
+  });
+
   describe('getSelectedAccount', () => {
     it.each([
       {

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.
    *
@@ -1305,6 +1321,11 @@ export class AccountsController extends BaseController<
       this.getAccount.bind(this),
     );
 
+    this.messagingSystem.registerActionHandler(
+      `AccountsController:getAccounts`,
+      this.getAccounts.bind(this),
+    );
+
     this.messagingSystem.registerActionHandler(
       `AccountsController:updateAccountMetadata`,
       this.updateAccountMetadata.bind(this),

packages/accounts-controller/src/index.ts

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

packages/keyring-controller/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add actions for `createNewVaultAndKeychain` and `createNewVaultAndRestore` ([#6654](https://github.com/MetaMask/core/pull/6708))
+  - These actions are meant to to be consumed by the `MultichainAccountService` in its `createMultichainAccountWallet` method.
+
 ### Changed
 
 - Bump `@metamask/utils` from `^11.4.2` to `^11.8.1` ([#6588](https://github.com/MetaMask/core/pull/6588), [#6708](https://github.com/MetaMask/core/pull/6708))

packages/keyring-controller/src/KeyringController.ts

@@ -184,6 +184,16 @@ export type KeyringControllerWithKeyringAction = {
   handler: KeyringController['withKeyring'];
 };
 
+export type KeyringControllerCreateNewVaultAndKeychainAction = {
+  type: `${typeof name}:createNewVaultAndKeychain`;
+  handler: KeyringController['createNewVaultAndKeychain'];
+};
+
+export type KeyringControllerCreateNewVaultAndRestoreAction = {
+  type: `${typeof name}:createNewVaultAndRestore`;
+  handler: KeyringController['createNewVaultAndRestore'];
+};
+
 export type KeyringControllerAddNewKeyringAction = {
   type: `${typeof name}:addNewKeyring`;
   handler: KeyringController['addNewKeyring'];
@@ -226,7 +236,9 @@ export type KeyringControllerActions =
   | KeyringControllerSignUserOperationAction
   | KeyringControllerAddNewAccountAction
   | KeyringControllerWithKeyringAction
-  | KeyringControllerAddNewKeyringAction;
+  | KeyringControllerAddNewKeyringAction
+  | KeyringControllerCreateNewVaultAndKeychainAction
+  | KeyringControllerCreateNewVaultAndRestoreAction;
 
 export type KeyringControllerEvents =
   | KeyringControllerStateChangeEvent
@@ -1780,6 +1792,16 @@ export class KeyringController extends BaseController<
       `${name}:addNewKeyring`,
       this.addNewKeyring.bind(this),
     );
+
+    this.messagingSystem.registerActionHandler(
+      `${name}:createNewVaultAndKeychain`,
+      this.createNewVaultAndKeychain.bind(this),
+    );
+
+    this.messagingSystem.registerActionHandler(
+      `${name}:createNewVaultAndRestore`,
+      this.createNewVaultAndRestore.bind(this),
+    );
   }
 
   /**

packages/multichain-account-service/CHANGELOG.md

@@ -10,10 +10,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Add `{Btc/Trx}AccountProvider` account providers ([#6662](https://github.com/MetaMask/core/pull/6662))
+- **BREAKING** A performance refactor was made around all the classes in this package ([#6654](https://github.com/MetaMask/core/pull/6708))
+  - Add logic in the `createMultichainAccountWallet` method in `MultichainAccountService` so that it can handle all entry points: importing an SRP, recovering a vault and creating a new vault.
+  - Add a `getAccountIds` method which returns all the account ids pertaining to a group.
+  - Add an `addAccounts` method on the `BaseBip44AccountProvider` class which keeps track of all the account IDs that pertain to it.
 
 ### Changed
 
 - Bump `@metamask/utils` from `^11.8.0` to `^11.8.1` ([#6708](https://github.com/MetaMask/core/pull/6708))
+- **BREAKING** A performance refactor was made around all the classes in this package ([#6654](https://github.com/MetaMask/core/pull/6708))
+  - The `MultichainAccountService` is refactored to construct a top level service state for its `init` function, this state is passed down to the `MultichainAccountWallet` and `MultichainAccountGroup` classes in slices for them to construct their internal states.
+  - Additional state is generated at the entry points where it needs to be updated i.e. `createMultichainAccountGroup`, `discoverAccounts` and `alignAccounts`.
+  - We no longer prevent group creation if some providers' `createAccounts` calls fail during group creation, only if they all fail.
+  - The `getAccounts` method in the `BaseBip44AccountProvider` class no longer relies on fetching the entire list of internal accounts from the `AccountsController`, instead it gets the specific accounts that it stores in its internal accounts list.
+  - The `EvmAccountProvider` no longer fetches from the `AccountController` to get an account for its ID, we deterministically get the associated account ID through `getUUIDFromAddressOfNormalAccount`.
+  - The `EvmAccountProvider` now uses the `getAccount` method from the `AccountsController` when fetching an account after account creation as it is more efficient.
+
+### Removed
+
+- **BREAKING** A performance refactor was made around all the classes in this package ([#6654](https://github.com/MetaMask/core/pull/6708))
+  - Remove `#handleOnAccountAdded` and `#handleOnAccountRemoved` methods in `MultichainAccountService` due to internal state being updated within the service.
+  - Remove `getAccountContext` (and associated map) in the `MultichainAccountService` as the service no longer uses that method.
+  - Remove the `sync` method in favor of the sole `init` method for both `MultichainAccountWallet` and `MultichainAccountGroup`.
 
 ## [1.2.0]
 

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

@@ -9,8 +9,12 @@ import type { Messenger } from '@metamask/base-controller';
 import { EthScope, SolScope } from '@metamask/keyring-api';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 
-import { MultichainAccountGroup } from './MultichainAccountGroup';
+import {
+  type GroupState,
+  MultichainAccountGroup,
+} from './MultichainAccountGroup';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
+import type { BaseBip44AccountProvider } from './providers';
 import type { MockAccountProvider } from './tests';
 import {
   MOCK_SNAP_ACCOUNT_2,
@@ -54,23 +58,40 @@ function setup({
   group: MultichainAccountGroup<Bip44Account<InternalAccount>>;
   providers: MockAccountProvider[];
 } {
-  const providers = accounts.map((providerAccounts) => {
-    return setupNamedAccountProvider({ accounts: providerAccounts });
+  const providers = accounts.map((providerAccounts, idx) => {
+    return setupNamedAccountProvider({
+      name: `Provider ${idx + 1}`,
+      accounts: providerAccounts,
+    });
   });
 
   const wallet = new MultichainAccountWallet<Bip44Account<InternalAccount>>({
     entropySource: MOCK_WALLET_1_ENTROPY_SOURCE,
     messenger: getMultichainAccountServiceMessenger(messenger),
-    providers,
+    providers: providers as unknown as BaseBip44AccountProvider[],
   });
 
   const group = new MultichainAccountGroup({
     wallet,
     groupIndex,
-    providers,
+    providers: providers as unknown as BaseBip44AccountProvider[],
     messenger: getMultichainAccountServiceMessenger(messenger),
   });
 
+  // Initialize group state from provided accounts so that constructor tests
+  // observe accounts immediately
+  const groupState = providers.reduce<GroupState>((state, provider, idx) => {
+    const ids = accounts[idx]
+      .filter((a) => 'options' in a && a.options?.entropy)
+      .map((a) => a.id);
+    if (ids.length > 0) {
+      state[provider.getName()] = ids;
+    }
+    return state;
+  }, {});
+
+  group.init(groupState);
+
   return { wallet, group, providers };
 }
 
@@ -95,6 +116,10 @@ describe('MultichainAccount', () => {
       expect(group.type).toBe(AccountGroupType.MultichainAccount);
       expect(group.groupIndex).toBe(groupIndex);
       expect(group.wallet).toStrictEqual(wallet);
+      expect(group.hasAccounts()).toBe(true);
+      expect(group.getAccountIds()).toStrictEqual(
+        expectedAccounts.map((a) => a.id),
+      );
       expect(group.getAccounts()).toHaveLength(expectedAccounts.length);
       expect(group.getAccounts()).toStrictEqual(expectedAccounts);
     });
@@ -177,6 +202,10 @@ describe('MultichainAccount', () => {
         ],
       });
 
+      providers[1].createAccounts.mockResolvedValueOnce([
+        MOCK_WALLET_1_SOL_ACCOUNT,
+      ]);
+
       await group.alignAccounts();
 
       expect(providers[0].createAccounts).not.toHaveBeenCalled();
@@ -199,7 +228,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,
@@ -219,7 +248,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- 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('Unable to create accounts'),
+      );
+
+      providers[2].createAccounts.mockRejectedValueOnce(
+        new Error('Unable to create accounts'),
+      );
+
+      await group.alignAccounts();
+
+      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- Provider 2: Unable to create accounts\n- Provider 3: Unable to create accounts`,
       );
     });
   });