From 486b9e9081d473a28326ff52ed7fa0a2253fbd1a Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 3 Nov 2025 13:02:12 +0100 Subject: [PATCH 1/8] fix(account-tree-controller)!: automatically removes "outdated" accounts from groups --- .../src/AccountTreeController.test.ts | 51 ++++++- .../src/AccountTreeController.ts | 126 +++++++++++++----- packages/account-tree-controller/src/types.ts | 8 +- .../tests/mockMessenger.ts | 1 + 4 files changed, 151 insertions(+), 35 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index d98b819b97d..6d87636ba79 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -1,4 +1,8 @@ -import type { AccountWalletId, Bip44Account } from '@metamask/account-api'; +import type { + AccountWalletId, + Bip44Account, + MultichainAccountGroup, +} from '@metamask/account-api'; import { AccountGroupType, AccountWalletType, @@ -10,6 +14,7 @@ import { } from '@metamask/account-api'; import type { AccountId } from '@metamask/accounts-controller'; import { deriveStateFromMetadata } from '@metamask/base-controller'; +import type { KeyringAccount } from '@metamask/keyring-api'; import { BtcAccountType, EthAccountType, @@ -1340,6 +1345,50 @@ describe('AccountTreeController', () => { }); }); + describe('on MultichainAccountService:multichainAccountGroupUpdated', () => { + it('removes account from the tree if their no longer part of the multichain account group', () => { + const { controller, messenger } = setup({ + accounts: [MOCK_HD_ACCOUNT_1, MOCK_TRX_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + messenger.registerActionHandler( + 'SnapController:get', + () => + // TODO: Update this to avoid the unknown cast if possible. + MOCK_SNAP_3 as unknown as ReturnType< + SnapControllerGetSnap['handler'] + >, + ); + + controller.init(); + + const walletId = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, + ); + const groupId = toMultichainAccountGroupId(walletId, 0); + + let group = controller.getAccountGroupObject(groupId); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(2); + + messenger.publish( + 'MultichainAccountService:multichainAccountGroupUpdated', + // Partial implementation, so we need to cast here. + { + id: groupId, + // Simulate a group with just 1 account now, meaning that `MOCK_TRX_ACCOUNT_1` got + // "removed/disabled" from this multichain account group. + getAccounts: jest.fn().mockReturnValue([MOCK_HD_ACCOUNT_1]), + } as unknown as MultichainAccountGroup>, + ); + + group = controller.getAccountGroupObject(groupId); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(1); + }); + }); + describe('account ordering by type', () => { it('orders accounts in group according to ACCOUNT_TYPE_TO_SORT_ORDER regardless of insertion order', () => { const evmAccount = MOCK_HD_ACCOUNT_1; diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index 2180076e233..dce7319aeeb 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -5,12 +5,15 @@ import type { AccountSelector, MultichainAccountWalletId, AccountGroupType, + Bip44Account, } from '@metamask/account-api'; import type { MultichainAccountWalletStatus } from '@metamask/account-api'; +import type { MultichainAccountGroup } from '@metamask/account-api'; import { type AccountId } from '@metamask/accounts-controller'; import type { StateMetadata } from '@metamask/base-controller'; import { BaseController } from '@metamask/base-controller'; import type { TraceCallback } from '@metamask/controller-utils'; +import type { KeyringAccount } from '@metamask/keyring-api'; import { isEvmAccountType } from '@metamask/keyring-api'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import { assert } from '@metamask/utils'; @@ -247,6 +250,13 @@ export class AccountTreeController extends BaseController< }, ); + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupUpdated', + (group) => { + this.#handleMultichainAccountWalletGroupUpdated(group); + }, + ); + this.#registerMessageHandlers(); } @@ -846,43 +856,63 @@ export class AccountTreeController extends BaseController< return; } - const context = this.#accountIdToContext.get(accountId); + this.#removeAccounts([accountId]); + } + + /** + * Removes an account from the tree if it exists. If not, this method does nothing. + * + * @param accountIds - Account IDs to remove from the tree. + */ + #removeAccounts(accountIds: AccountId[]) { + const removedAccounts: AccountId[] = []; + + const previousSelectedAccountGroup = + this.state.accountTree.selectedAccountGroup; + let selectedAccountGroupChanged = false; - if (context) { - const { walletId, groupId } = context; + this.update((state) => { + for (const accountId of accountIds) { + const context = this.#accountIdToContext.get(accountId); - const previousSelectedAccountGroup = - this.state.accountTree.selectedAccountGroup; - let selectedAccountGroupChanged = false; + if (context) { + const { walletId, groupId } = context; - this.update((state) => { - const accounts = - state.accountTree.wallets[walletId]?.groups[groupId]?.accounts; - - if (accounts) { - const index = accounts.indexOf(accountId); - if (index !== -1) { - accounts.splice(index, 1); - - // Check if we need to update selectedAccountGroup after removal - if ( - state.accountTree.selectedAccountGroup === groupId && - accounts.length === 0 - ) { - // The currently selected group is now empty, find a new group to select - const newSelectedAccountGroup = this.#getDefaultAccountGroupId( - state.accountTree.wallets, - ); - state.accountTree.selectedAccountGroup = newSelectedAccountGroup; - selectedAccountGroupChanged = - newSelectedAccountGroup !== previousSelectedAccountGroup; + const accounts = + state.accountTree.wallets[walletId]?.groups[groupId]?.accounts; + + if (accounts) { + const index = accounts.indexOf(accountId); + if (index !== -1) { + accounts.splice(index, 1); + + // Now we know this account got removed. + removedAccounts.push(accountId); + + // Check if we need to update selectedAccountGroup after removal + if ( + state.accountTree.selectedAccountGroup === groupId && + accounts.length === 0 + ) { + // The currently selected group is now empty, find a new group to select + const newSelectedAccountGroup = this.#getDefaultAccountGroupId( + state.accountTree.wallets, + ); + state.accountTree.selectedAccountGroup = + newSelectedAccountGroup; + selectedAccountGroupChanged = + newSelectedAccountGroup !== previousSelectedAccountGroup; + } + } + if (accounts.length === 0) { + this.#pruneEmptyGroupAndWallet(state, walletId, groupId); } - } - if (accounts.length === 0) { - this.#pruneEmptyGroupAndWallet(state, walletId, groupId); } } - }); + } + }); + + if (removedAccounts.length) { this.messenger.publish( `${controllerName}:accountTreeChange`, this.state.accountTree, @@ -898,7 +928,9 @@ export class AccountTreeController extends BaseController< } // Clear reverse-mapping for that account. - this.#accountIdToContext.delete(accountId); + removedAccounts.forEach((accountId) => + this.#accountIdToContext.delete(accountId), + ); } } @@ -1219,6 +1251,36 @@ export class AccountTreeController extends BaseController< }); } + /** + * Handles multichain account group updates. + * + * @param group - Multichain account group that got updated. + */ + #handleMultichainAccountWalletGroupUpdated( + group: MultichainAccountGroup>, + ): void { + const multichainGroupAccounts = new Set( + group.getAccounts().map((account) => account.id), + ); + + // We only check if some accounts are no longer part of the multichain + // account group. This is required mainly with the `AccountProviderWrapper`s that + // can be disabled when "Basic functionality" is turned OFF or when the wrappers + // are controlled by remote feature flags. + const treeGroup = this.#getAccountGroup(group.id); + const treeGroupAccountsToRemove: AccountId[] = []; + if (treeGroup) { + for (const account of treeGroup.accounts) { + // Remove accounts that not longer exist on the multichain account group. + if (!multichainGroupAccounts.has(account)) { + treeGroupAccountsToRemove.push(account); + } + } + } + + this.#removeAccounts(treeGroupAccountsToRemove); + } + /** * Gets account group object. * diff --git a/packages/account-tree-controller/src/types.ts b/packages/account-tree-controller/src/types.ts index da288a1b565..83f324614d4 100644 --- a/packages/account-tree-controller/src/types.ts +++ b/packages/account-tree-controller/src/types.ts @@ -39,7 +39,10 @@ import type { AccountWalletObject, AccountTreeWalletPersistedMetadata, } from './wallet'; -import type { MultichainAccountServiceWalletStatusChangeEvent } from '../../multichain-account-service/src/types'; +import type { + MultichainAccountServiceMultichainAccountGroupUpdatedEvent, + MultichainAccountServiceWalletStatusChangeEvent, +} from '../../multichain-account-service/src/types'; // Backward compatibility aliases using indexed access types /** @@ -167,7 +170,8 @@ export type AllowedEvents = | AccountsControllerAccountRemovedEvent | AccountsControllerSelectedAccountChangeEvent | UserStorageController.UserStorageControllerStateChangeEvent - | MultichainAccountServiceWalletStatusChangeEvent; + | MultichainAccountServiceWalletStatusChangeEvent + | MultichainAccountServiceMultichainAccountGroupUpdatedEvent; export type AccountTreeControllerEvents = | AccountTreeControllerStateChangeEvent diff --git a/packages/account-tree-controller/tests/mockMessenger.ts b/packages/account-tree-controller/tests/mockMessenger.ts index 10fa770a675..203aa1097c3 100644 --- a/packages/account-tree-controller/tests/mockMessenger.ts +++ b/packages/account-tree-controller/tests/mockMessenger.ts @@ -50,6 +50,7 @@ export function getAccountTreeControllerMessenger( 'AccountsController:selectedAccountChange', 'UserStorageController:stateChange', 'MultichainAccountService:walletStatusChange', + 'MultichainAccountService:multichainAccountGroupUpdated', ], actions: [ 'AccountsController:listMultichainAccounts', From c222de5b640d0873819982ffa39f47acbb4a7398 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 3 Nov 2025 13:09:05 +0100 Subject: [PATCH 2/8] chore: update changelog --- packages/account-tree-controller/CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md index b09828e5624..2836f700038 100644 --- a/packages/account-tree-controller/CHANGELOG.md +++ b/packages/account-tree-controller/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- **BREAKING:** Properly removes "outdated" accounts from the `MultichainAccountService` ([#7038](https://github.com/MetaMask/core/pull/7038)) + - The `MultichainAccountService:multichainAccountGroupUpdated` event is now required. + - This change will ensure that "disabled" accounts are not properly removed from the account tree too. + ## [2.0.0] ### Changed From 60778d929eb8379d12e45dfffda173e60f8ce061 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 3 Nov 2025 15:54:48 +0100 Subject: [PATCH 3/8] fix: also re-sync account with the service during init --- .../src/AccountTreeController.test.ts | 14 +++++++ .../src/AccountTreeController.ts | 41 +++++++++++++++++-- packages/account-tree-controller/src/types.ts | 6 ++- .../tests/mockMessenger.ts | 2 + 4 files changed, 58 insertions(+), 5 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index 6d87636ba79..514de79e4a9 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -332,6 +332,9 @@ function setup({ getAccount: jest.fn(), getSelectedMultichainAccount: jest.fn(), }, + MultichainAccountService: { + getMultichainAccountWallets: jest.fn(), + }, UserStorageController: { getState: jest.fn(), performGetStorage: jest.fn(), @@ -428,6 +431,17 @@ function setup({ ); } + // Using an empty list of wallets will make sure that no multichain accounts got + // removed automatically removed from the tree upon calling `init`. That's what + // we want for almost all tests. + mocks.MultichainAccountService.getMultichainAccountWallets.mockReturnValue( + [], + ); + messenger.registerActionHandler( + 'MultichainAccountService:getMultichainAccountWallets', + mocks.MultichainAccountService.getMultichainAccountWallets, + ); + const accountTreeControllerMessenger = getAccountTreeControllerMessenger(messenger); const controller = new AccountTreeController({ diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index dce7319aeeb..107af63decf 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -250,10 +250,17 @@ export class AccountTreeController extends BaseController< }, ); + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupCreated', + (group) => { + this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group); + }, + ); + this.messenger.subscribe( 'MultichainAccountService:multichainAccountGroupUpdated', (group) => { - this.#handleMultichainAccountWalletGroupUpdated(group); + this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group); }, ); @@ -307,6 +314,17 @@ export class AccountTreeController extends BaseController< this.#insert(wallets, account); } + // Since we're building the tree out of of the account list, we don't know if those + // accounts got disabled at the service level. To make sure both states are "in-sync" + // we check each multichain account groups and remove extra accounts if there's any! + for (const wallet of this.messenger.call( + 'MultichainAccountService:getMultichainAccountWallets', + )) { + for (const group of wallet.getMultichainAccountGroups()) { + this.#resyncAccountsFromMultichainAccountGroups(group); + } + } + // Once we have the account tree, we can apply persisted metadata (names + UI states). let previousSelectedAccountGroupStillExists = false; this.update((state) => { @@ -1252,11 +1270,15 @@ export class AccountTreeController extends BaseController< } /** - * Handles multichain account group updates. + * Check if an account group object is in-sync with its multichain + * account group counterpart. * - * @param group - Multichain account group that got updated. + * If not, then the extra-accounts fromt he account group object + * will be removed from the tree. + * + * @param group - Multichain account group that got created or updated. */ - #handleMultichainAccountWalletGroupUpdated( + #resyncAccountsFromMultichainAccountGroups( group: MultichainAccountGroup>, ): void { const multichainGroupAccounts = new Set( @@ -1281,6 +1303,17 @@ export class AccountTreeController extends BaseController< this.#removeAccounts(treeGroupAccountsToRemove); } + /** + * Handles multichain account group updates. + * + * @param group - Multichain account group that got updated. + */ + #handleMultichainAccountWalletGroupCreatedOrUpdated( + group: MultichainAccountGroup>, + ): void { + this.#resyncAccountsFromMultichainAccountGroups(group); + } + /** * Gets account group object. * diff --git a/packages/account-tree-controller/src/types.ts b/packages/account-tree-controller/src/types.ts index 83f324614d4..89e38afc911 100644 --- a/packages/account-tree-controller/src/types.ts +++ b/packages/account-tree-controller/src/types.ts @@ -40,6 +40,8 @@ import type { AccountTreeWalletPersistedMetadata, } from './wallet'; import type { + MultichainAccountServiceGetMultichainAccountWalletsAction, + MultichainAccountServiceMultichainAccountGroupCreatedEvent, MultichainAccountServiceMultichainAccountGroupUpdatedEvent, MultichainAccountServiceWalletStatusChangeEvent, } from '../../multichain-account-service/src/types'; @@ -130,7 +132,8 @@ export type AllowedActions = | UserStorageController.UserStorageControllerPerformSetStorage | UserStorageController.UserStorageControllerPerformBatchSetStorage | AuthenticationController.AuthenticationControllerGetSessionProfile - | MultichainAccountServiceCreateMultichainAccountGroupAction; + | MultichainAccountServiceCreateMultichainAccountGroupAction + | MultichainAccountServiceGetMultichainAccountWalletsAction; export type AccountTreeControllerActions = | AccountTreeControllerGetStateAction @@ -171,6 +174,7 @@ export type AllowedEvents = | AccountsControllerSelectedAccountChangeEvent | UserStorageController.UserStorageControllerStateChangeEvent | MultichainAccountServiceWalletStatusChangeEvent + | MultichainAccountServiceMultichainAccountGroupCreatedEvent | MultichainAccountServiceMultichainAccountGroupUpdatedEvent; export type AccountTreeControllerEvents = diff --git a/packages/account-tree-controller/tests/mockMessenger.ts b/packages/account-tree-controller/tests/mockMessenger.ts index 203aa1097c3..4dfa8081d3f 100644 --- a/packages/account-tree-controller/tests/mockMessenger.ts +++ b/packages/account-tree-controller/tests/mockMessenger.ts @@ -50,6 +50,7 @@ export function getAccountTreeControllerMessenger( 'AccountsController:selectedAccountChange', 'UserStorageController:stateChange', 'MultichainAccountService:walletStatusChange', + 'MultichainAccountService:multichainAccountGroupCreated', 'MultichainAccountService:multichainAccountGroupUpdated', ], actions: [ @@ -64,6 +65,7 @@ export function getAccountTreeControllerMessenger( 'UserStorageController:performBatchSetStorage', 'AuthenticationController:getSessionProfile', 'MultichainAccountService:createMultichainAccountGroup', + 'MultichainAccountService:getMultichainAccountWallets', 'KeyringController:getState', 'SnapController:get', ], From 4cd03fb5a52842cea029bee6403c41b916cb7cb1 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 3 Nov 2025 17:37:40 +0100 Subject: [PATCH 4/8] fix: properly remove extra accounts during init --- .../src/AccountTreeController.ts | 39 +++++++++++++------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index 107af63decf..6b07483a280 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -321,7 +321,15 @@ export class AccountTreeController extends BaseController< 'MultichainAccountService:getMultichainAccountWallets', )) { for (const group of wallet.getMultichainAccountGroups()) { - this.#resyncAccountsFromMultichainAccountGroups(group); + const { inSync } = + this.#getDiffAccountsFromMultichainAccountGroups(group); + + const groupObject = wallets[wallet.id]?.groups[group.id]; + if (groupObject && inSync.length > 0) { + // We still check for `inSync` length to make sure we can cast here (since groups + // MUST HAVE at least 1 account). + groupObject.accounts = inSync as AccountGroupObject['accounts']; + } } } @@ -1270,37 +1278,42 @@ export class AccountTreeController extends BaseController< } /** - * Check if an account group object is in-sync with its multichain + * Gets the difference between an account group object and its multichain * account group counterpart. * - * If not, then the extra-accounts fromt he account group object - * will be removed from the tree. - * * @param group - Multichain account group that got created or updated. + * @returns The diff object to know which accounts are similar and the ones that + * needs to be removed. */ - #resyncAccountsFromMultichainAccountGroups( + #getDiffAccountsFromMultichainAccountGroups( group: MultichainAccountGroup>, - ): void { + ): { inSync: AccountId[]; toRemove: AccountId[] } { const multichainGroupAccounts = new Set( group.getAccounts().map((account) => account.id), ); + const sync = { + inSync: [] as AccountId[], + toRemove: [] as AccountId[], + }; + // We only check if some accounts are no longer part of the multichain // account group. This is required mainly with the `AccountProviderWrapper`s that // can be disabled when "Basic functionality" is turned OFF or when the wrappers // are controlled by remote feature flags. const treeGroup = this.#getAccountGroup(group.id); - const treeGroupAccountsToRemove: AccountId[] = []; if (treeGroup) { for (const account of treeGroup.accounts) { // Remove accounts that not longer exist on the multichain account group. - if (!multichainGroupAccounts.has(account)) { - treeGroupAccountsToRemove.push(account); + if (multichainGroupAccounts.has(account)) { + sync.inSync.push(account); + } else { + sync.toRemove.push(account); } } } - this.#removeAccounts(treeGroupAccountsToRemove); + return sync; } /** @@ -1311,7 +1324,9 @@ export class AccountTreeController extends BaseController< #handleMultichainAccountWalletGroupCreatedOrUpdated( group: MultichainAccountGroup>, ): void { - this.#resyncAccountsFromMultichainAccountGroups(group); + const { toRemove } = + this.#getDiffAccountsFromMultichainAccountGroups(group); + this.#removeAccounts(toRemove); } /** From f86a6bc4cf8eca6ac4b88d1e44e9e1bd4bb7e96e Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Tue, 4 Nov 2025 14:02:21 +0100 Subject: [PATCH 5/8] refactor: refactor --- .../src/AccountTreeController.test.ts | 80 +++++++++--- .../src/AccountTreeController.ts | 117 ++++++++++-------- 2 files changed, 131 insertions(+), 66 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index 514de79e4a9..74a6f45eb0f 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -6,6 +6,7 @@ import type { import { AccountGroupType, AccountWalletType, + isBip44Account, toAccountGroupId, toAccountWalletId, toMultichainAccountGroupId, @@ -14,7 +15,7 @@ import { } from '@metamask/account-api'; import type { AccountId } from '@metamask/accounts-controller'; import { deriveStateFromMetadata } from '@metamask/base-controller'; -import type { KeyringAccount } from '@metamask/keyring-api'; +import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api'; import { BtcAccountType, EthAccountType, @@ -237,6 +238,48 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = { const mockGetSelectedMultichainAccountActionHandler = jest.fn(); +/** + * Get multichain account wallets for the given list of accounts. + * + * @param accounts Accounts to build wallets for. + * @returns Multichain account wallets. + */ +function getMultichainAccountWalletsFromAccounts(accounts: KeyringAccount[]) { + const bip44Accounts = accounts.filter(isBip44Account); + + const wallets: { + [entropySource: EntropySourceId]: { + [groupIndex: number]: Set>; + }; + } = {}; + for (const bip44Account of bip44Accounts) { + const { id: entropySource, groupIndex } = bip44Account.options.entropy; + + wallets[entropySource] ??= {}; + wallets[entropySource][groupIndex] ??= new Set(); + wallets[entropySource][groupIndex].add(bip44Account); + } + + return Object.entries(wallets).map(([entropySource, groups]) => { + const walletId = toMultichainAccountWalletId(entropySource); + return { + id: walletId, + getMultichainAccountGroups: jest.fn().mockImplementation(() => { + return Object.entries(groups).map(([groupIndex, groupAccounts]) => { + const groupId = toMultichainAccountGroupId( + walletId, + Number(groupIndex), + ); + return { + id: groupId, + getAccounts: jest.fn().mockReturnValue(Array.from(groupAccounts)), + }; + }); + }), + }; + }); +} + /** * Sets up the AccountTreeController for testing. * @@ -309,6 +352,9 @@ function setup({ getSelectedMultichainAccount: jest.Mock; getAccount: jest.Mock; }; + MultichainAccountService: { + getMultichainAccountWallets: jest.Mock; + }; UserStorageController: { performGetStorage: jest.Mock; performGetStorageAllFeatureEntries: jest.Mock; @@ -418,6 +464,17 @@ function setup({ 'UserStorageController:performBatchSetStorage', mocks.UserStorageController.performBatchSetStorage, ); + + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => + getMultichainAccountWalletsFromAccounts( + mocks.AccountsController.accounts, + ), + ); + messenger.registerActionHandler( + 'MultichainAccountService:getMultichainAccountWallets', + mocks.MultichainAccountService.getMultichainAccountWallets, + ); } if (keyrings) { @@ -431,17 +488,6 @@ function setup({ ); } - // Using an empty list of wallets will make sure that no multichain accounts got - // removed automatically removed from the tree upon calling `init`. That's what - // we want for almost all tests. - mocks.MultichainAccountService.getMultichainAccountWallets.mockReturnValue( - [], - ); - messenger.registerActionHandler( - 'MultichainAccountService:getMultichainAccountWallets', - mocks.MultichainAccountService.getMultichainAccountWallets, - ); - const accountTreeControllerMessenger = getAccountTreeControllerMessenger(messenger); const controller = new AccountTreeController({ @@ -962,12 +1008,18 @@ describe('AccountTreeController', () => { }); const now = Date.now(); - mocks.AccountsController.listMultichainAccounts.mockReturnValue([ + const mockAccounts = [ // Faking accounts to be out of order: mockAccountWith(1, now + 1000), mockAccountWith(2, now + 2000), mockAccountWith(0, now), - ]); + ]; + mocks.AccountsController.listMultichainAccounts.mockReturnValue( + mockAccounts, + ); + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => getMultichainAccountWalletsFromAccounts(mockAccounts), + ); controller.init(); diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index 6b07483a280..dc2bdeca87a 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -1,4 +1,8 @@ -import { AccountWalletType, select } from '@metamask/account-api'; +import { + AccountWalletType, + isBip44Account, + select, +} from '@metamask/account-api'; import type { AccountGroupId, AccountWalletId, @@ -309,28 +313,23 @@ export class AccountTreeController extends BaseController< (a, b) => a.metadata.importTime - b.metadata.importTime, ); + // HACK: Since we're not consuming accounts directly from the service, we have to + // cross-check with its internal state before inserting initial BIP-44 accounts. + // Those accounts could be marked as "disabled" (using the + // `AccountProviderWrapper`) and if they are, we don't want to insert them in the tree. + // The real fix for this would be consume accounts from the service instead of the + // `AccountsController`, but this requires a bit more testing, for now we keep this + // simple. + const bip44Accounts = this.#getBip44AccountIds(); + // For now, we always re-compute all wallets, we do not re-use the existing state. for (const account of accounts) { - this.#insert(wallets, account); - } - - // Since we're building the tree out of of the account list, we don't know if those - // accounts got disabled at the service level. To make sure both states are "in-sync" - // we check each multichain account groups and remove extra accounts if there's any! - for (const wallet of this.messenger.call( - 'MultichainAccountService:getMultichainAccountWallets', - )) { - for (const group of wallet.getMultichainAccountGroups()) { - const { inSync } = - this.#getDiffAccountsFromMultichainAccountGroups(group); - - const groupObject = wallets[wallet.id]?.groups[group.id]; - if (groupObject && inSync.length > 0) { - // We still check for `inSync` length to make sure we can cast here (since groups - // MUST HAVE at least 1 account). - groupObject.accounts = inSync as AccountGroupObject['accounts']; - } + if (isBip44Account(account) && !bip44Accounts.has(account.id)) { + // We skip those, so they won't be part of the initial tree. + continue; } + + this.#insert(wallets, account); } // Once we have the account tree, we can apply persisted metadata (names + UI states). @@ -1278,55 +1277,69 @@ export class AccountTreeController extends BaseController< } /** - * Gets the difference between an account group object and its multichain - * account group counterpart. + * Get account IDs from a multichain account group. * - * @param group - Multichain account group that got created or updated. - * @returns The diff object to know which accounts are similar and the ones that - * needs to be removed. + * @param group - Multichain account group to extract account IDs from. + * @returns Set of BIP-44 account IDs for the given group. */ - #getDiffAccountsFromMultichainAccountGroups( + #getBip44AccountIdsFromMultichainAccountGroup( group: MultichainAccountGroup>, - ): { inSync: AccountId[]; toRemove: AccountId[] } { - const multichainGroupAccounts = new Set( - group.getAccounts().map((account) => account.id), - ); + ): Set['id']> { + // FIXME: We should introduce a way to just get the account IDs since getting + // account object might be a bit more costly (at least with the current + // architecture). + return new Set(group.getAccounts().map((account) => account.id)); + } - const sync = { - inSync: [] as AccountId[], - toRemove: [] as AccountId[], - }; + /** + * Get account IDs owned by the multichain account service. + * + * @returns Set of BIP-44 account IDs. + */ + #getBip44AccountIds() { + const accounts = new Set(); - // We only check if some accounts are no longer part of the multichain - // account group. This is required mainly with the `AccountProviderWrapper`s that - // can be disabled when "Basic functionality" is turned OFF or when the wrappers - // are controlled by remote feature flags. - const treeGroup = this.#getAccountGroup(group.id); - if (treeGroup) { - for (const account of treeGroup.accounts) { - // Remove accounts that not longer exist on the multichain account group. - if (multichainGroupAccounts.has(account)) { - sync.inSync.push(account); - } else { - sync.toRemove.push(account); + // Since we're building the tree out of of the account list, we don't know if those + // accounts got disabled at the service level. To make sure both states are "in-sync" + // we check each multichain account groups and remove extra accounts if there's any! + for (const wallet of this.messenger.call( + 'MultichainAccountService:getMultichainAccountWallets', + )) { + for (const group of wallet.getMultichainAccountGroups()) { + for (const account of this.#getBip44AccountIdsFromMultichainAccountGroup( + group, + )) { + accounts.add(account); } } } - return sync; + return accounts; } /** * Handles multichain account group updates. * - * @param group - Multichain account group that got updated. + * @param multichainGroup - Multichain account group that got updated. */ #handleMultichainAccountWalletGroupCreatedOrUpdated( - group: MultichainAccountGroup>, + multichainGroup: MultichainAccountGroup>, ): void { - const { toRemove } = - this.#getDiffAccountsFromMultichainAccountGroups(group); - this.#removeAccounts(toRemove); + const bip44GroupAccounts = + this.#getBip44AccountIdsFromMultichainAccountGroup(multichainGroup); + + const group = this.#getAccountGroup(multichainGroup.id); + if (group) { + this.#removeAccounts( + group.accounts.filter( + // We only check if some accounts are no longer part of the multichain + // account group. This is required mainly with the `AccountProviderWrapper`s that + // can be disabled when "Basic functionality" is turned OFF or when the wrappers + // are controlled by remote feature flags. + (account) => !bip44GroupAccounts.has(account), + ), + ); + } } /** From 08cbb81abe2800308248cee1972ad8ae0fdaab3c Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Tue, 4 Nov 2025 14:56:25 +0100 Subject: [PATCH 6/8] chore: update changelog --- packages/account-tree-controller/CHANGELOG.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md index 2836f700038..02ccc548bad 100644 --- a/packages/account-tree-controller/CHANGELOG.md +++ b/packages/account-tree-controller/CHANGELOG.md @@ -10,8 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - **BREAKING:** Properly removes "outdated" accounts from the `MultichainAccountService` ([#7038](https://github.com/MetaMask/core/pull/7038)) - - The `MultichainAccountService:multichainAccountGroupUpdated` event is now required. - This change will ensure that "disabled" accounts are not properly removed from the account tree too. + - The `MultichainAccountService:multichainAccountGroupCreated` event is now required. + - The `MultichainAccountService:multichainAccountGroupUpdated` event is now required. + - The `MultichainAccountService:getMultichainAccountWallets` action is now required. ## [2.0.0] From ceb67ff46acbead4a89528f203d9594acdde73e6 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Tue, 4 Nov 2025 15:02:49 +0100 Subject: [PATCH 7/8] test: add test for :multichainAccountGroupCreated --- .../src/AccountTreeController.test.ts | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index 74a6f45eb0f..7b48d63c726 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -1455,6 +1455,39 @@ describe('AccountTreeController', () => { }); }); + describe('on MultichainAccountService:multichainAccountGroupCreated', () => { + it('does nothing if an multichain account group got created before inserting any accounts', () => { + const { controller, messenger } = setup({ + accounts: [], + keyrings: [MOCK_HD_KEYRING_1], + }); + + controller.init(); + + const walletId = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, + ); + const groupId = toMultichainAccountGroupId(walletId, 0); + + let group = controller.getAccountGroupObject(groupId); + expect(group).toBeUndefined(); + + messenger.publish( + 'MultichainAccountService:multichainAccountGroupCreated', + // Partial implementation, so we need to cast here. + { + id: groupId, + getAccounts: jest.fn().mockReturnValue([MOCK_HD_ACCOUNT_1]), + } as unknown as MultichainAccountGroup>, + ); + + // Still undefined, since we're only reacting to `:accountAdded` for now. + group = controller.getAccountGroupObject(groupId); + expect(group).toBeUndefined(); + }); + }); + + describe('account ordering by type', () => { it('orders accounts in group according to ACCOUNT_TYPE_TO_SORT_ORDER regardless of insertion order', () => { const evmAccount = MOCK_HD_ACCOUNT_1; From 3c46556c4a2d645b3dc474dde061aa544ab8c728 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Tue, 4 Nov 2025 15:12:34 +0100 Subject: [PATCH 8/8] test: add test for init --- .../src/AccountTreeController.test.ts | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index 7b48d63c726..fe7b055e078 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -1035,6 +1035,30 @@ describe('AccountTreeController', () => { expect(groupIds[1]).toBe(toMultichainAccountGroupId(walletId, 1)); expect(groupIds[2]).toBe(toMultichainAccountGroupId(walletId, 2)); }); + + it('skips disabled bip-44 accounts', () => { + const { controller, mocks } = setup({ + accounts: [MOCK_HD_ACCOUNT_1, MOCK_TRX_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + // We simulate TRX account to be "disabled" here, so the `AccountsController` has 2 accounts (EVM + TRX), but + // our wallets will only returns the EVM account here, meaning TRX will be skipped. + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => getMultichainAccountWalletsFromAccounts([MOCK_HD_ACCOUNT_1]), + ); + + controller.init(); + + const group = controller.getAccountGroupObject( + toMultichainAccountGroupId( + toMultichainAccountWalletId(MOCK_HD_KEYRING_1.metadata.id), + MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, + ), + ); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(1); // Just EVM, no TRX. + }); }); describe('getAccountGroupObject', () => { @@ -1487,7 +1511,6 @@ describe('AccountTreeController', () => { }); }); - describe('account ordering by type', () => { it('orders accounts in group according to ACCOUNT_TYPE_TO_SORT_ORDER regardless of insertion order', () => { const evmAccount = MOCK_HD_ACCOUNT_1;