fix(multichain-account-service): prevent spurrious alignment logs + smaller map updates

Author: ccharly

packages/multichain-account-service/CHANGELOG.md

@@ -10,6 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Bump `@metamask/accounts-controller` from `^36.0.0` to `^36.0.1` ([#7996](https://github.com/MetaMask/core/pull/7996))
+- Faster state change during alignment ([#8136](https://github.com/MetaMask/core/pull/8136))
+  - We used to re-create the entire mapping of accounts and their providers during alignment, now we only update the accounts that actually changed or got removed.
+
+### Fixed
+
+- Prevent spurious logs during alignment ([#8136](https://github.com/MetaMask/core/pull/8136))
 
 ## [7.0.0]
 

packages/multichain-account-service/jest.config.js

@@ -14,6 +14,9 @@ module.exports = merge(baseConfig, {
   // The display name when running multiple projects
   displayName,
 
+  // Exclude test helpers from coverage collection
+  coveragePathIgnorePatterns: ['.*/src/tests/.*'],
+
   // An object that configures minimum threshold enforcement for coverage results
   coverageThreshold: {
     global: {

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

@@ -283,6 +283,24 @@ describe('MultichainAccount', () => {
       expect(providers[1].createAccounts).not.toHaveBeenCalled();
     });
 
+    it('removes accounts from the group when a provider no longer returns them during alignment', async () => {
+      const groupIndex = 0;
+      const { group, providers } = setup({
+        groupIndex,
+        accounts: [[MOCK_WALLET_1_EVM_ACCOUNT, MOCK_WALLET_1_SOL_ACCOUNT], []],
+      });
+
+      // Provider 0 previously had both accounts but now only returns the EVM one.
+      mockCreateAccountsOnce(providers[0], [MOCK_WALLET_1_EVM_ACCOUNT]);
+
+      await group.alignAccounts();
+
+      expect(group.getAccount(MOCK_WALLET_1_EVM_ACCOUNT.id)).toBe(
+        MOCK_WALLET_1_EVM_ACCOUNT,
+      );
+      expect(group.getAccount(MOCK_WALLET_1_SOL_ACCOUNT.id)).toBeUndefined();
+    });
+
     it('removes accounts from the group when a provider is disabled', async () => {
       const groupIndex = 0;
       const { group, providers } = setup({

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

@@ -267,33 +267,71 @@ export class MultichainAccountGroup<
   async alignAccounts(): Promise<void> {
     this.#log('Aligning accounts...');
 
-    this.#providerToAccounts.clear();
-    this.#accountToProvider.clear();
-
     const results = await Promise.allSettled(
       this.#providers.map(async (provider) => {
+        const providerName = provider.getName();
+
         try {
+          // Existing (old) accounts for that provider (if any).
+          const oldAccounts = new Set(
+            this.#providerToAccounts.get(provider) ?? [],
+          );
+
+          // Remove any previously tracked accounts if provider gets disabled.
+          if (isAccountProviderWrapper(provider) && provider.isDisabled()) {
+            this.#log(
+              `Account provider "${providerName}" is disabled, skipping alignment...`,
+            );
+
+            for (const accountId of oldAccounts) {
+              this.#accountToProvider.delete(accountId);
+            }
+            this.#providerToAccounts.delete(provider);
+
+            return [];
+          }
+
+          // We align accounts no matter what, we cannot guess if providers starts to support new
+          // account types or scopes.
           const accounts = await provider.alignAccounts({
             entropySource: this.wallet.entropySource,
             groupIndex: this.groupIndex,
           });
 
-          const isDisabled =
-            isAccountProviderWrapper(provider) && provider.isDisabled();
+          // Compute a diff between the previously known accounts to see if some accounts got removed
+          // or added during alignment.
+          const newAccounts = new Set(accounts);
+          for (const accountId of accounts) {
+            // If we knew about this account before, it's not new nor removed.
+            if (oldAccounts.has(accountId)) {
+              oldAccounts.delete(accountId);
+              newAccounts.delete(accountId);
+            }
+          }
 
-          if (isDisabled) {
+          const hasRemovedAccounts = oldAccounts.size > 0;
+          const hasAddedAccounts = newAccounts.size > 0;
+          if (hasRemovedAccounts) {
             this.#log(
-              `Account provider "${provider.getName()}" is disabled, skipping alignment...`,
+              `Found ${oldAccounts.size} removed accounts for account provider "${providerName}", removing them from the group...`,
             );
-          } else if (accounts.length > 0) {
+
+            for (const accountId of oldAccounts) {
+              this.#accountToProvider.delete(accountId);
+            }
+          }
+          if (hasAddedAccounts) {
             this.#log(
-              `Found missing accounts for account provider "${provider.getName()}", creating them now...`,
+              `Found ${newAccounts.size} new accounts for account provider "${providerName}", adding them to the group...`,
             );
-            this.#providerToAccounts.set(provider, accounts);
-            for (const accountId of accounts) {
+
+            for (const accountId of newAccounts) {
               this.#accountToProvider.set(accountId, provider);
             }
           }
+          if (hasRemovedAccounts || hasAddedAccounts) {
+            this.#providerToAccounts.set(provider, accounts);
+          }
 
           return accounts;
         } catch (error) {
@@ -302,11 +340,11 @@ export class MultichainAccountGroup<
             `${WARNING_PREFIX} ${error instanceof Error ? error.message : String(error)}`,
           );
           const sentryError = createSentryError(
-            `Unable to align accounts with provider "${provider.getName()}"`,
+            `Unable to align accounts with provider "${providerName}"`,
             error as Error,
             {
               groupIndex: this.groupIndex,
-              provider: provider.getName(),
+              provider: providerName,
             },
           );
           this.#messenger.captureException?.(sentryError);

packages/multichain-account-service/src/providers/AccountProviderWrapper.test.ts

@@ -0,0 +1,118 @@
+import type { Bip44Account } from '@metamask/account-api';
+import type {
+  CreateAccountOptions,
+  KeyringAccount,
+  KeyringCapabilities,
+} from '@metamask/keyring-api';
+import type { InternalAccount } from '@metamask/keyring-internal-api';
+
+import { AccountProviderWrapper } from './AccountProviderWrapper';
+import { BaseBip44AccountProvider } from './BaseBip44AccountProvider';
+import {
+  getMultichainAccountServiceMessenger,
+  getRootMessenger,
+  MOCK_HD_KEYRING_1,
+  MOCK_SOL_ACCOUNT_1,
+} from '../tests';
+import type { MultichainAccountServiceMessenger } from '../types';
+
+class MockInnerProvider extends BaseBip44AccountProvider {
+  readonly capabilities: KeyringCapabilities = {
+    scopes: [],
+    bip44: { deriveIndex: true },
+  };
+
+  getName(): string {
+    return 'MockInnerProvider';
+  }
+
+  isAccountCompatible(_account: Bip44Account<KeyringAccount>): boolean {
+    return true;
+  }
+
+  async createAccounts(
+    _options: CreateAccountOptions,
+  ): Promise<Bip44Account<KeyringAccount>[]> {
+    return [];
+  }
+
+  async discoverAccounts(_options: {
+    entropySource: string;
+    groupIndex: number;
+  }): Promise<Bip44Account<KeyringAccount>[]> {
+    return [];
+  }
+
+  resyncAccounts(_accounts: Bip44Account<InternalAccount>[]): Promise<void> {
+    return Promise.resolve();
+  }
+}
+
+function setup(): {
+  wrapper: AccountProviderWrapper;
+  innerProvider: MockInnerProvider;
+} {
+  const messenger = getRootMessenger();
+  const multichainMessenger: MultichainAccountServiceMessenger =
+    getMultichainAccountServiceMessenger(messenger);
+
+  const innerProvider = new MockInnerProvider(multichainMessenger);
+  const wrapper = new AccountProviderWrapper(multichainMessenger, innerProvider);
+
+  return { wrapper, innerProvider };
+}
+
+describe('AccountProviderWrapper', () => {
+  describe('alignAccounts', () => {
+    it('returns empty array when disabled', async () => {
+      const { wrapper } = setup();
+
+      wrapper.setEnabled(false);
+
+      const result = await wrapper.alignAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(result).toStrictEqual([]);
+    });
+
+    it('does not delegate to inner provider when disabled', async () => {
+      const { wrapper, innerProvider } = setup();
+      innerProvider.init([MOCK_SOL_ACCOUNT_1.id]);
+
+      const alignSpy = jest
+        .spyOn(innerProvider, 'alignAccounts')
+        .mockResolvedValue([MOCK_SOL_ACCOUNT_1.id]);
+
+      wrapper.setEnabled(false);
+
+      await wrapper.alignAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(alignSpy).not.toHaveBeenCalled();
+    });
+
+    it('delegates to inner provider when enabled', async () => {
+      const { wrapper, innerProvider } = setup();
+      innerProvider.init([MOCK_SOL_ACCOUNT_1.id]);
+
+      const alignSpy = jest
+        .spyOn(innerProvider, 'alignAccounts')
+        .mockResolvedValue([MOCK_SOL_ACCOUNT_1.id]);
+
+      const result = await wrapper.alignAccounts({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+
+      expect(alignSpy).toHaveBeenCalledWith({
+        entropySource: MOCK_HD_KEYRING_1.metadata.id,
+        groupIndex: 0,
+      });
+      expect(result).toStrictEqual([MOCK_SOL_ACCOUNT_1.id]);
+    });
+  });
+});