feat(multichain-account-service): basic functionality management

packages/multichain-account-service/CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - 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.
+- Add `setBasicFunctionality` method to control providers state and trigger wallets alignment ([#6332](https://github.com/MetaMask/core/pull/6332))
+  - Add `ProviderWrapper` to handle Snap account providers behavior when disabled
 
 ### Changed
 

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

@@ -7,6 +7,7 @@ import { KeyringTypes, type KeyringObject } from '@metamask/keyring-controller';
 
 import { MultichainAccountService } from './MultichainAccountService';
 import { EvmAccountProvider } from './providers/EvmAccountProvider';
+import { ProviderWrapper } from './providers/ProviderWrapper';
 import { SolAccountProvider } from './providers/SolAccountProvider';
 import type { MockAccountProvider } from './tests';
 import {
@@ -626,6 +627,59 @@ describe('MultichainAccountService', () => {
     });
   });
 
+  describe('getIsAlignmentInProgress', () => {
+    it('returns false initially', () => {
+      const { service } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+      });
+      expect(service.getIsAlignmentInProgress()).toBe(false);
+    });
+
+    it('returns true during alignWallets and false after completion', async () => {
+      const { service } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+      });
+
+      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({
+        accounts: [MOCK_HD_ACCOUNT_1],
+      });
+
+      const alignmentPromise = service.alignWallet(
+        MOCK_HD_KEYRING_1.metadata.id,
+      );
+      expect(service.getIsAlignmentInProgress()).toBe(true);
+
+      await alignmentPromise;
+      expect(service.getIsAlignmentInProgress()).toBe(false);
+    });
+
+    it('returns false after alignment error', async () => {
+      const { service, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+      });
+
+      mocks.EvmAccountProvider.createAccounts.mockRejectedValueOnce(
+        new Error('Test error'),
+      );
+
+      try {
+        await service.alignWallets();
+      } catch {
+        // Expected to throw
+      }
+
+      expect(service.getIsAlignmentInProgress()).toBe(false);
+    });
+  });
+
   describe('actions', () => {
     it('gets a multichain account with MultichainAccountService:getMultichainAccount', () => {
       const accounts = [MOCK_HD_ACCOUNT_1];
@@ -749,5 +803,154 @@ describe('MultichainAccountService', () => {
         groupIndex: 0,
       });
     });
+
+    it('sets basic functionality with MultichainAccountService:setBasicFunctionality', async () => {
+      const { messenger } = setup({ accounts: [MOCK_HD_ACCOUNT_1] });
+
+      // This tests the action handler registration
+      expect(
+        await messenger.call(
+          'MultichainAccountService:setBasicFunctionality',
+          true,
+        ),
+      ).toBeUndefined();
+      expect(
+        await messenger.call(
+          'MultichainAccountService:setBasicFunctionality',
+          false,
+        ),
+      ).toBeUndefined();
+    });
+
+    it('gets alignment progress with MultichainAccountService:getIsAlignmentInProgress', () => {
+      const { messenger } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+      });
+
+      const isInProgress = messenger.call(
+        'MultichainAccountService:getIsAlignmentInProgress',
+      );
+
+      expect(isInProgress).toBe(false);
+    });
+  });
+
+  describe('setBasicFunctionality', () => {
+    it('accepts a boolean parameter instead of object', async () => {
+      const { service } = setup({ accounts: [MOCK_HD_ACCOUNT_1] });
+
+      // These should not throw errors
+      expect(await service.setBasicFunctionality(true)).toBeUndefined();
+      expect(await service.setBasicFunctionality(false)).toBeUndefined();
+    });
+
+    it('can be called with boolean true', async () => {
+      const { service } = setup({ accounts: [MOCK_HD_ACCOUNT_1] });
+
+      // This tests the simplified parameter signature
+      expect(await service.setBasicFunctionality(true)).toBeUndefined();
+    });
+
+    it('can be called with boolean false', async () => {
+      const { service } = setup({ accounts: [MOCK_HD_ACCOUNT_1] });
+
+      // This tests the simplified parameter signature
+      expect(await service.setBasicFunctionality(false)).toBeUndefined();
+    });
+  });
+
+  describe('ProviderWrapper disabled behavior', () => {
+    let mockProvider: MockAccountProvider;
+    let wrapper: ProviderWrapper;
+
+    beforeEach(() => {
+      const { mocks } = setup({ accounts: [MOCK_HD_ACCOUNT_1] });
+      mockProvider = mocks.SolAccountProvider;
+      wrapper = new ProviderWrapper(mockProvider);
+    });
+
+    it('returns empty array when getAccounts() is disabled', () => {
+      // Enable first - should work normally
+      mockProvider.getAccounts.mockReturnValue([MOCK_HD_ACCOUNT_1]);
+      expect(wrapper.getAccounts()).toStrictEqual([MOCK_HD_ACCOUNT_1]);
+
+      // Disable - should return empty array
+      wrapper.setEnabled(false);
+      expect(wrapper.getAccounts()).toStrictEqual([]);
+    });
+
+    it('throws error when getAccount() is disabled', () => {
+      // Enable first - should work normally
+      mockProvider.getAccount.mockReturnValue(MOCK_HD_ACCOUNT_1);
+      expect(wrapper.getAccount('test-id')).toStrictEqual(MOCK_HD_ACCOUNT_1);
+
+      // Disable - should throw error
+      wrapper.setEnabled(false);
+      expect(() => wrapper.getAccount('test-id')).toThrow('is disabled');
+    });
+
+    it('returns empty array when createAccounts() is disabled', async () => {
+      const options = {
+        entropySource: MOCK_HD_ACCOUNT_1.options.entropy.id,
+        groupIndex: 0,
+      };
+
+      // Enable first - should work normally
+      mockProvider.createAccounts.mockResolvedValue([MOCK_HD_ACCOUNT_1]);
+      expect(await wrapper.createAccounts(options)).toStrictEqual([
+        MOCK_HD_ACCOUNT_1,
+      ]);
+
+      // Disable - should return empty array and not call underlying provider
+      wrapper.setEnabled(false);
+
+      const result = await wrapper.createAccounts(options);
+      expect(result).toStrictEqual([]);
+    });
+
+    it('returns empty array when discoverAndCreateAccounts() is disabled', async () => {
+      const options = {
+        entropySource: MOCK_HD_ACCOUNT_1.options.entropy.id,
+        groupIndex: 0,
+      };
+
+      // Enable first - should work normally
+      mockProvider.discoverAndCreateAccounts.mockResolvedValue([
+        MOCK_HD_ACCOUNT_1,
+      ]);
+      expect(await wrapper.discoverAndCreateAccounts(options)).toStrictEqual([
+        MOCK_HD_ACCOUNT_1,
+      ]);
+
+      // Disable - should return empty array
+      wrapper.setEnabled(false);
+
+      const result = await wrapper.discoverAndCreateAccounts(options);
+      expect(result).toStrictEqual([]);
+    });
+
+    it('proxies isAccountCompatible() correctly', () => {
+      // Test when provider has the method
+      const providerWithMethod = mockProvider as MockAccountProvider & {
+        isAccountCompatible: jest.Mock;
+      };
+      jest
+        .spyOn(providerWithMethod, 'isAccountCompatible')
+        .mockImplementation()
+        .mockReturnValue(true);
+      expect(wrapper.isAccountCompatible(MOCK_HD_ACCOUNT_1)).toBe(true);
+      expect(providerWithMethod.isAccountCompatible).toHaveBeenCalledWith(
+        MOCK_HD_ACCOUNT_1,
+      );
+
+      // Test when provider doesn't have the method (fallback to true)
+      const providerWithoutMethod = { ...mockProvider };
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (providerWithoutMethod as any).isAccountCompatible = undefined;
+      const wrapperWithoutMethod = new ProviderWrapper(providerWithoutMethod);
+      expect(wrapperWithoutMethod.isAccountCompatible(MOCK_HD_ACCOUNT_1)).toBe(
+        true,
+      );
+    });
   });
 });

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

@@ -5,14 +5,18 @@ import {
 import type {
   MultichainAccountWalletId,
   Bip44Account,
+  AccountProvider,
 } from '@metamask/account-api';
-import type { AccountProvider } from '@metamask/account-api';
 import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api';
 import { KeyringTypes } from '@metamask/keyring-controller';
 
 import type { MultichainAccountGroup } from './MultichainAccountGroup';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
 import { EvmAccountProvider } from './providers/EvmAccountProvider';
+import {
+  ProviderWrapper,
+  isProviderWrapper,
+} from './providers/ProviderWrapper';
 import { SolAccountProvider } from './providers/SolAccountProvider';
 import type { MultichainAccountServiceMessenger } from './types';
 
@@ -40,7 +44,10 @@ type AccountContext<Account extends Bip44Account<KeyringAccount>> = {
 export class MultichainAccountService {
   readonly #messenger: MultichainAccountServiceMessenger;
 
-  readonly #providers: AccountProvider<Bip44Account<KeyringAccount>>[];
+  readonly #providers: (
+    | AccountProvider<Bip44Account<KeyringAccount>>
+    | ProviderWrapper
+  )[];
 
   readonly #wallets: Map<
     MultichainAccountWalletId,
@@ -73,10 +80,11 @@ export class MultichainAccountService {
     this.#messenger = messenger;
     this.#wallets = new Map();
     this.#accountIdToContext = new Map();
+
     // TODO: Rely on keyring capabilities once the keyring API is used by all keyrings.
     this.#providers = [
       new EvmAccountProvider(this.#messenger),
-      new SolAccountProvider(this.#messenger),
+      new ProviderWrapper(new SolAccountProvider(this.#messenger)),
       // Custom account providers that can be provided by the MetaMask client.
       ...providers,
     ];
@@ -105,6 +113,10 @@ export class MultichainAccountService {
       'MultichainAccountService:createMultichainAccountGroup',
       (...args) => this.createMultichainAccountGroup(...args),
     );
+    this.#messenger.registerActionHandler(
+      'MultichainAccountService:setBasicFunctionality',
+      (...args) => this.setBasicFunctionality(...args),
+    );
     this.#messenger.registerActionHandler(
       'MultichainAccountService:alignWallets',
       (...args) => this.alignWallets(...args),
@@ -113,6 +125,10 @@ export class MultichainAccountService {
       'MultichainAccountService:alignWallet',
       (...args) => this.alignWallet(...args),
     );
+    this.#messenger.registerActionHandler(
+      'MultichainAccountService:getIsAlignmentInProgress',
+      () => this.getIsAlignmentInProgress(),
+    );
   }
 
   /**
@@ -359,6 +375,40 @@ export class MultichainAccountService {
     );
   }
 
+  /**
+   * Set basic functionality state and trigger alignment if enabled.
+   * When basic functionality is disabled, snap-based providers are disabled.
+   * When enabled, all snap providers are enabled and wallet alignment is triggered.
+   * EVM providers are never disabled as they're required for basic wallet functionality.
+   *
+   * @param enabled - Whether basic functionality is enabled.
+   */
+  async setBasicFunctionality(enabled: boolean): Promise<void> {
+    // Loop through providers and enable/disable only wrapped ones when basic functionality changes
+    for (const provider of this.#providers) {
+      if (isProviderWrapper(provider)) {
+        provider.setEnabled(enabled);
+      }
+      // Regular providers (like EVM) are never disabled for basic functionality
+    }
+
+    // Trigger alignment only when basic functionality is enabled
+    if (enabled) {
+      await this.alignWallets();
+    }
+  }
+
+  /**
+   * Gets whether wallet alignment is currently in progress.
+   *
+   * @returns True if any wallet alignment is in progress, false otherwise.
+   */
+  getIsAlignmentInProgress(): boolean {
+    return Array.from(this.#wallets.values()).some((wallet) =>
+      wallet.getIsAlignmentInProgress(),
+    );
+  }
+
   /**
    * Align all multichain account wallets.
    */