From 5e1e04b55443c5b3880832d04434c1f29c815967 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 13 Apr 2026 13:35:57 +0200 Subject: [PATCH 1/3] feat(account-tree-controller): persist accountTree --- .../src/AccountTreeController.test.ts | 77 +++++++++++++++++++ .../src/AccountTreeController.ts | 51 +++++++++++- 2 files changed, 127 insertions(+), 1 deletion(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index ae6a5e1e16d..c7e01154309 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -235,6 +235,46 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = { const mockGetSelectedMultichainAccountActionHandler = jest.fn(); +const MOCK_PREPOPULATED_WALLET_ID = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, +); +const MOCK_PREPOPULATED_GROUP_ID = toMultichainAccountGroupId( + MOCK_PREPOPULATED_WALLET_ID, + MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, +); +const MOCK_PREPOPULATED_STATE: Partial = { + selectedAccountGroup: MOCK_PREPOPULATED_GROUP_ID, + accountTree: { + wallets: { + [MOCK_PREPOPULATED_WALLET_ID]: { + id: MOCK_PREPOPULATED_WALLET_ID, + type: AccountWalletType.Entropy, + status: 'ready', + groups: { + [MOCK_PREPOPULATED_GROUP_ID]: { + id: MOCK_PREPOPULATED_GROUP_ID, + type: AccountGroupType.MultichainAccount, + accounts: [MOCK_HD_ACCOUNT_1.id], + metadata: { + name: 'Account 1', + entropy: { + groupIndex: MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, + }, + pinned: false, + hidden: false, + lastSelected: 0, + }, + }, + }, + metadata: { + name: 'Wallet 1', + entropy: { id: MOCK_HD_KEYRING_1.metadata.id }, + }, + }, + }, + }, +}; + /** * Sets up the AccountTreeController for testing. * @@ -1076,6 +1116,27 @@ describe('AccountTreeController', () => { toMultichainAccountWalletId(MOCK_HD_KEYRING_2.metadata.id), ); }); + + it('populates reverse mappings from persisted tree state without calling init()', () => { + const { controller } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + state: MOCK_PREPOPULATED_STATE, + }); + + // init() is NOT called — reverse mappings must be populated from the persisted tree. + const context = controller.getAccountContext(MOCK_HD_ACCOUNT_1.id); + expect(context).toBeDefined(); + expect(context?.walletId).toBe(MOCK_PREPOPULATED_WALLET_ID); + expect(context?.groupId).toBe(MOCK_PREPOPULATED_GROUP_ID); + expect(context?.sortOrder).toBeDefined(); + + const group = controller.getAccountGroupObject( + MOCK_PREPOPULATED_GROUP_ID, + ); + expect(group).toBeDefined(); + expect(group?.id).toBe(MOCK_PREPOPULATED_GROUP_ID); + }); }); describe('getAccountsFromSelectAccountGroup', () => { @@ -1147,6 +1208,19 @@ describe('AccountTreeController', () => { expect(controller.getAccountsFromSelectedAccountGroup()).toHaveLength(0); }); + + it('returns accounts from persisted state without calling init()', () => { + const { controller } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + state: MOCK_PREPOPULATED_STATE, + }); + + // init() is NOT called — accounts must be served from the persisted state. + expect(controller.getAccountsFromSelectedAccountGroup()).toStrictEqual([ + MOCK_HD_ACCOUNT_1, + ]); + }); }); describe('on AccountsController:accountsRemoved', () => { @@ -4863,6 +4937,9 @@ describe('AccountTreeController', () => { ).toMatchInlineSnapshot(` { "accountGroupsMetadata": {}, + "accountTree": { + "wallets": {}, + }, "accountWalletsMetadata": {}, "hasAccountTreeSyncingSyncedAtLeastOnce": false, "selectedAccountGroup": "", diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index f1839d969aa..daaa22f6657 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -65,7 +65,7 @@ const accountTreeControllerMetadata: StateMetadata = { accountTree: { includeInStateLogs: true, - persist: false, // We do re-recompute this state everytime. + persist: true, includeInDebugSnapshot: false, usedInUi: true, }, @@ -235,6 +235,12 @@ export class AccountTreeController extends BaseController< this.#createBackupAndSyncContext(), ); + // We build the initial tree (from the state), but this might get updated + // after `init` got called. Having it available now, allow our consumers to + // re-use the state that was already available before we restart/lock the + // wallet. + this.#initTreeContext(this.state.accountTree); + this.messenger.subscribe('AccountsController:accountsAdded', (accounts) => { this.#handleAccountsAdded(accounts); }); @@ -275,6 +281,49 @@ export class AccountTreeController extends BaseController< ); } + /** + * Initialize the account tree from the given state. + * + * @param tree The account tree state to initialize from. + */ + #initTreeContext(tree: AccountTreeControllerState['accountTree']): void { + for (const [walletId, wallet] of Object.entries(tree.wallets) as [ + AccountWalletId, + AccountWalletObject, + ][]) { + for (const [groupId, group] of Object.entries(wallet.groups) as [ + AccountGroupId, + AccountGroupObject, + ][]) { + this.#groupIdToWalletId.set(groupId, walletId); + + // We still need to go through accounts for the sort order. + for (const accountId of group.accounts) { + // The `AccountsController` also persists its accounts, so we + // can fetch them too! + const account = this.messenger.call( + 'AccountsController:getAccount', + accountId, + ); + + // NOTE: The account should always be available, but if it is not, we + // still let it inside the tree with a max sort order to avoid blocking + // the entire tree construction. + let sortOrder = MAX_SORT_ORDER; + if (account) { + sortOrder = ACCOUNT_TYPE_TO_SORT_ORDER[account.type]; + } + + this.#accountIdToContext.set(accountId, { + walletId, + groupId, + sortOrder, + }); + } + } + } + } + /** * Initialize the controller's state. * From 8c20fbfb8508c36d05ece95260ddb3b472c7644d Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Mon, 13 Apr 2026 18:00:56 +0200 Subject: [PATCH 2/3] chore: changelog --- packages/account-tree-controller/CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md index 01c09285810..3f43e00799a 100644 --- a/packages/account-tree-controller/CHANGELOG.md +++ b/packages/account-tree-controller/CHANGELOG.md @@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- Now persist `accounTree` ([#8437](https://github.com/MetaMask/core/pull/8437)) + - The tree is not persisted and can be used before `init` is called. + - This allow consumers (like the assets controllers) to rely on the selected group's accounts right away. - Remove dynamic identifiers (wallet IDs, group IDs) from backup and sync thrown error messages to improve Sentry error grouping ([#8349](https://github.com/MetaMask/core/pull/8349)) - Bump `@metamask/accounts-controller` from `^37.1.1` to `^37.2.0` ([#8363](https://github.com/MetaMask/core/pull/8363)) - Bump `@metamask/keyring-controller` from `^25.1.1` to `^25.2.0` ([#8363](https://github.com/MetaMask/core/pull/8363)) From 65dedea5f0e4c83537780870f0be5d5b84e9a662 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Tue, 14 Apr 2026 15:12:22 +0200 Subject: [PATCH 3/3] refactor: use listMultichainAccounts instead of N getAccount --- .../src/AccountTreeController.test.ts | 10 ++++++++++ .../src/AccountTreeController.ts | 12 ++++++------ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index c7e01154309..9e9aa08089b 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -917,6 +917,11 @@ describe('AccountTreeController', () => { keyrings: [MOCK_HD_KEYRING_1], }); + // Reset mock call counts after construction: the constructor calls + // #initTreeContext which seeds the fast-index Maps from persisted state, + // incrementing the call counters before we even call init(). + jest.clearAllMocks(); + controller.init(); expect( mocks.AccountsController.listMultichainAccounts, @@ -942,6 +947,11 @@ describe('AccountTreeController', () => { keyrings: [MOCK_HD_KEYRING_1], }); + // Reset mock call counts after construction: the constructor calls + // #initTreeContext which seeds the fast-index Maps from persisted state, + // incrementing the call counters before we even call init(). + jest.clearAllMocks(); + controller.init(); expect( mocks.AccountsController.listMultichainAccounts, diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index daaa22f6657..5093de0217c 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -287,6 +287,11 @@ export class AccountTreeController extends BaseController< * @param tree The account tree state to initialize from. */ #initTreeContext(tree: AccountTreeControllerState['accountTree']): void { + // Fetch persisted accounts from the `AccountsController`. + const accounts = new Map( + this.#listAccounts().map((account) => [account.id, account]), + ); + for (const [walletId, wallet] of Object.entries(tree.wallets) as [ AccountWalletId, AccountWalletObject, @@ -299,12 +304,7 @@ export class AccountTreeController extends BaseController< // We still need to go through accounts for the sort order. for (const accountId of group.accounts) { - // The `AccountsController` also persists its accounts, so we - // can fetch them too! - const account = this.messenger.call( - 'AccountsController:getAccount', - accountId, - ); + const account = accounts.get(accountId); // NOTE: The account should always be available, but if it is not, we // still let it inside the tree with a max sort order to avoid blocking