feat(account-tree-controller)!: add :removeAccountWallet action

packages/account-tree-controller/CHANGELOG.md

@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **BREAKING:** Add `AccountTreeController:removeAccountWallet`([#10253](https://github.com/MetaMask/core/pull/10253))
+  - It removes all accounts in a non-primary wallet.
+  - It rejects removal of the primary HD wallet.
+  - We now need those new actions on the messenger `MultichainAccountService:removeMultichainAccountWallet` and `KeyringController:removeAccount`.
+
 ### Changed
 
 - Bump `@metamask/profile-sync-controller` from `^32.0.0` to `^32.1.0` ([#10184](https://github.com/MetaMask/core/pull/10184))

packages/account-tree-controller/src/AccountTreeController-method-action-types.ts

@@ -39,6 +39,25 @@ export type AccountTreeControllerIsInitializedAction = {
   handler: AccountTreeController['isInitialized'];
 };
 
+/**
+ * Removes an account wallet and all of its underlying accounts.
+ *
+ * The account tree is a derived view of AccountsController state, so this
+ * method intentionally does not mutate tree nodes directly. Account removal
+ * causes AccountsController to publish `accountsRemoved`, which lets
+ * `#handleAccountsRemoved` consistently prune tree nodes, reverse mappings,
+ * metadata, and selection state.
+ *
+ * @param walletId - Account wallet ID.
+ * @throws If the account tree has not been initialized.
+ * @throws If the wallet does not exist.
+ * @throws If the wallet belongs to the primary HD keyring.
+ */
+export type AccountTreeControllerRemoveAccountWalletAction = {
+  type: `AccountTreeController:removeAccountWallet`;
+  handler: AccountTreeController['removeAccountWallet'];
+};
+
 /**
  * Gets the account wallet object from its ID.
  *
@@ -289,6 +308,7 @@ export type AccountTreeControllerMethodActions =
   | AccountTreeControllerInitAction
   | AccountTreeControllerReinitAction
   | AccountTreeControllerIsInitializedAction
+  | AccountTreeControllerRemoveAccountWalletAction
   | AccountTreeControllerGetAccountWalletObjectAction
   | AccountTreeControllerGetAccountWalletObjectsAction
   | AccountTreeControllerGetAccountsFromSelectedAccountGroupAction

packages/account-tree-controller/src/AccountTreeController.test.ts

@@ -319,26 +319,30 @@
     consoleWarn: jest.SpyInstance;
   };
   mocks: {
     KeyringController: {
       keyrings: KeyringObject[];
       getState: jest.Mock;
+      removeAccount: jest.Mock;
       verifyPassword: jest.Mock;
       withController: jest.Mock;
     };
+    MultichainAccountService: {
+      removeMultichainAccountWallet: jest.Mock;
+    };
     AccountsController: {
       accounts: InternalAccount[];
       listMultichainAccounts: jest.Mock;
       getSelectedMultichainAccount: jest.Mock;
       getAccount: jest.Mock;
     };
     UserStorageController: {
       performGetStorage: jest.Mock;
       performGetStorageAllFeatureEntries: jest.Mock;
       performSetStorage: jest.Mock;
       performBatchSetStorage: jest.Mock;
       syncInternalAccountsWithUserStorage: jest.Mock;
     };
     AuthenticationController: {
       getSessionProfile: jest.Mock;
     };
   };
@@ -347,9 +351,13 @@
     KeyringController: {
       keyrings,
       getState: jest.fn(),
+      removeAccount: jest.fn().mockResolvedValue(undefined),
       verifyPassword: jest.fn().mockResolvedValue(undefined),
       withController: jest.fn(),
     },
+    MultichainAccountService: {
+      removeMultichainAccountWallet: jest.fn().mockResolvedValue(undefined),
+    },
     AccountsController: {
       accounts,
       listMultichainAccounts: jest.fn(),
@@ -456,6 +464,11 @@
       mocks.KeyringController.verifyPassword,
     );
 
+    messenger.registerActionHandler(
+      'KeyringController:removeAccount',
+      mocks.KeyringController.removeAccount,
+    );
+
     // Default: call the callback with no existing keyrings so private-key
     // imports are a no-op unless the test overrides this handler.
     mocks.KeyringController.withController.mockImplementation(
@@ -472,6 +485,11 @@
     );
   }
 
+  messenger.registerActionHandler(
+    'MultichainAccountService:removeMultichainAccountWallet',
+    mocks.MultichainAccountService.removeMultichainAccountWallet,
+  );
+
   const accountTreeControllerMessenger =
     getAccountTreeControllerMessenger(messenger);
   const controller = new AccountTreeController({
@@ -557,6 +575,229 @@
     });
   });
 
+  describe('removeAccountWallet', () => {
+    it('throws if the account tree is not initialized', async () => {
+      const { controller } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+        state: MOCK_PREPOPULATED_STATE,
+      });
+
+      await expect(
+        controller.removeAccountWallet(MOCK_PREPOPULATED_WALLET_ID),
+      ).rejects.toThrow('Account tree is not initialized');
+    });
+
+    it('throws if the wallet does not exist', async () => {
+      const { controller } = setup();
+      controller.init();
+
+      await expect(
+        controller.removeAccountWallet(
+          'entropy:missing-wallet' as AccountWalletId,
+        ),
+      ).rejects.toThrow('Account wallet not found in tree');
+    });
+
+    it('throws before removing the primary entropy wallet', async () => {
+      const { controller, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2],
+        keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
+      });
+      controller.init();
+
+      await expect(
+        controller.removeAccountWallet(
+          toMultichainAccountWalletId(MOCK_HD_KEYRING_1.metadata.id),
+        ),
+      ).rejects.toThrow('Cannot remove the primary account wallet');
+      expect(
+        mocks.MultichainAccountService.removeMultichainAccountWallet,
+      ).not.toHaveBeenCalled();
+      expect(mocks.KeyringController.removeAccount).not.toHaveBeenCalled();
+    });
+
+    it('removes a secondary entropy wallet through MultichainAccountService', async () => {
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2],
+        keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
+      });
+      controller.init();
+      const walletId = toMultichainAccountWalletId(
+        MOCK_HD_KEYRING_2.metadata.id,
+      );
+      mocks.MultichainAccountService.removeMultichainAccountWallet.mockImplementation(
+        async () => {
+          messenger.publish('AccountsController:accountsRemoved', [
+            MOCK_HD_ACCOUNT_2.id,
+          ]);
+        },
+      );
+
+      await controller.removeAccountWallet(walletId);
+
+      expect(
+        mocks.MultichainAccountService.removeMultichainAccountWallet,
+      ).toHaveBeenCalledWith(MOCK_HD_KEYRING_2.metadata.id);
+      expect(controller.getAccountWalletObject(walletId)).toBeUndefined();
+    });
+
+    it('force-removes every hardware account through KeyringController', async () => {
+      const secondHardwareAccount: InternalAccount = {
+        ...MOCK_HARDWARE_ACCOUNT_1,
+        id: 'mock-hardware-id-2',
+        address: '0xDEF',
+      };
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HARDWARE_ACCOUNT_1, secondHardwareAccount],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+      controller.init();
+      const walletId = toAccountWalletId(
+        AccountWalletType.Keyring,
+        KeyringTypes.ledger,
+      );
+      mocks.KeyringController.removeAccount.mockImplementation(
+        async (address) => {
+          const account = [MOCK_HARDWARE_ACCOUNT_1, secondHardwareAccount].find(
+            (candidate) => candidate.address === address,
+          );
+          messenger.publish('AccountsController:accountsRemoved', [
+            account?.id as AccountId,
+          ]);
+        },
+      );
+
+      await controller.removeAccountWallet(walletId);
+
+      expect(mocks.KeyringController.removeAccount).toHaveBeenCalledTimes(2);
+      expect(mocks.KeyringController.removeAccount).toHaveBeenNthCalledWith(
+        1,
+        MOCK_HARDWARE_ACCOUNT_1.address,
+      );
+      expect(mocks.KeyringController.removeAccount).toHaveBeenNthCalledWith(
+        2,
+        secondHardwareAccount.address,
+      );
+      expect(controller.getAccountWalletObject(walletId)).toBeUndefined();
+    });
+
+    it('force-removes Snap accounts through KeyringController', async () => {
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_SNAP_ACCOUNT_2],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+      messenger.registerActionHandler(
+        'SnapController:getSnap',
+        () =>
+          MOCK_SNAP_2 as unknown as ReturnType<
+            SnapControllerGetSnap['handler']
+          >,
+      );
+      controller.init();
+      const walletId = toAccountWalletId(
+        AccountWalletType.Snap,
+        MOCK_SNAP_2.id,
+      );
+      mocks.KeyringController.removeAccount.mockImplementation(async () => {
+        messenger.publish('AccountsController:accountsRemoved', [
+          MOCK_SNAP_ACCOUNT_2.id,
+        ]);
+      });
+
+      await controller.removeAccountWallet(walletId);
+
+      expect(mocks.KeyringController.removeAccount).toHaveBeenCalledWith(
+        MOCK_SNAP_ACCOUNT_2.address,
+      );
+      expect(controller.getAccountWalletObject(walletId)).toBeUndefined();
+    });
+
+    it('continues removing accounts after an individual removal fails', async () => {
+      const secondHardwareAccount: InternalAccount = {
+        ...MOCK_HARDWARE_ACCOUNT_1,
+        id: 'mock-hardware-id-2',
+        address: '0xDEF',
+      };
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HARDWARE_ACCOUNT_1, secondHardwareAccount],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+      controller.init();
+      const walletId = toAccountWalletId(
+        AccountWalletType.Keyring,
+        KeyringTypes.ledger,
+      );
+      mocks.KeyringController.removeAccount
+        .mockRejectedValueOnce(new Error('Removal failed'))
+        .mockImplementationOnce(async () => {
+          messenger.publish('AccountsController:accountsRemoved', [
+            secondHardwareAccount.id,
+          ]);
+        });
+
+      const consoleErrorSpy = jest
+        .spyOn(console, 'error')
+        .mockImplementation(() => undefined);
+
+      await controller.removeAccountWallet(walletId);
+
+      expect(mocks.KeyringController.removeAccount).toHaveBeenCalledTimes(2);
+      expect(
+        controller.getAccountWalletObject(walletId)?.groups[
+          toAccountGroupId(walletId, MOCK_HARDWARE_ACCOUNT_1.address)
+        ].accounts,
+      ).toStrictEqual([MOCK_HARDWARE_ACCOUNT_1.id]);
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'Account wallet removal is incomplete',
+        {
+          walletId,
+          remainingAccountIds: [MOCK_HARDWARE_ACCOUNT_1.id],
+        },
+      );
+    });
+
+    it('keeps the wallet when an account is missing from AccountsController', async () => {
+      const { controller, mocks } = setup({
+        accounts: [MOCK_HARDWARE_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+      controller.init();
+      const walletId = toAccountWalletId(
+        AccountWalletType.Keyring,
+        KeyringTypes.ledger,
+      );
+      mocks.AccountsController.accounts = [];
+      const consoleErrorSpy = jest
+        .spyOn(console, 'error')
+        .mockImplementation(() => undefined);
+
+      await controller.removeAccountWallet(walletId);
+
+      expect(mocks.KeyringController.removeAccount).not.toHaveBeenCalled();
+      expect(controller.getAccountWalletObject(walletId)).toBeDefined();
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'Account wallet removal is incomplete',
+        {
+          walletId,
+          remainingAccountIds: [MOCK_HARDWARE_ACCOUNT_1.id],
+        },
+      );
+    });
+
+    it('is exposed through the controller messenger', async () => {
+      const { accountTreeControllerMessenger, controller } = setup();
+      controller.init();
+
+      await expect(
+        accountTreeControllerMessenger.call(
+          'AccountTreeController:removeAccountWallet',
+          'entropy:missing-wallet' as AccountWalletId,
+        ),
+      ).rejects.toThrow('Account wallet not found in tree');
+    });
+  });
+
   describe('init', () => {
     it('groups accounts by entropy source, then snapId, then wallet type', () => {
       const { controller, messenger } = setup({