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)) diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index ae6a5e1e16d..9e9aa08089b 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. * @@ -877,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, @@ -902,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, @@ -1076,6 +1126,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 +1218,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 +4947,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..5093de0217c 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 { + // 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, + ][]) { + 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) { + 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 + // 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. *