diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md index b09828e5624..cf9bf49f223 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:** Use `MultichainAccountService` to build multichain account (BIP-44) nodes ([#6646](https://github.com/MetaMask/core/pull/6646)) + - Previously, the controller was using a similar matching logic for BIP-44 accounts, which was redundant with the logic from this service. + - Wallets and groups are now directly consumed from the service, making the service be the source of truth for accounts related to BIP-44. + ## [2.0.0] ### Changed diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index d98b819b97d..933c364023e 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -1,7 +1,13 @@ -import type { AccountWalletId, Bip44Account } from '@metamask/account-api'; +import type { + AccountWalletId, + Bip44Account, + MultichainAccountGroupId, + MultichainAccountWalletId, +} from '@metamask/account-api'; import { AccountGroupType, AccountWalletType, + isBip44Account, toAccountGroupId, toAccountWalletId, toMultichainAccountGroupId, @@ -10,6 +16,7 @@ import { } from '@metamask/account-api'; import type { AccountId } from '@metamask/accounts-controller'; import { deriveStateFromMetadata } from '@metamask/base-controller'; +import type { EntropySourceId } from '@metamask/keyring-api'; import { BtcAccountType, EthAccountType, @@ -26,6 +33,11 @@ import { import type { KeyringObject } from '@metamask/keyring-controller'; import { KeyringTypes } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { + MultichainAccountGroup, + MultichainAccountWallet, + MultichainAccountService, +} from '@metamask/multichain-account-service'; import type { GetSnap as SnapControllerGetSnap } from '@metamask/snaps-controllers'; import { @@ -107,6 +119,12 @@ const MOCK_HD_KEYRING_2 = { accounts: ['0x456'], }; +const MOCK_SIMPLE_KEYRING = { + type: KeyringTypes.simple, + metadata: { id: 'mock-simple-keyring-id', name: 'Simple Keyring' }, + accounts: ['0xabc', '0xdef'], +}; + const MOCK_HD_ACCOUNT_1: Bip44Account = { id: 'mock-id-1', address: '0x123', @@ -153,8 +171,8 @@ const MOCK_HD_ACCOUNT_2: Bip44Account = { }, }; -const MOCK_SNAP_ACCOUNT_1: Bip44Account = { - id: 'mock-snap-id-1', +const MOCK_SOL_ACCOUNT_1: Bip44Account = { + id: 'mock-sol-id-1', address: 'aabbccdd', options: { entropy: { @@ -176,6 +194,22 @@ const MOCK_SNAP_ACCOUNT_1: Bip44Account = { }, }; +const MOCK_SNAP_ACCOUNT_1: InternalAccount = { + id: 'mock-snap-id-1', + address: '0x789', + options: {}, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: '', + keyring: { type: KeyringTypes.snap }, + snap: MOCK_SNAP_1, + importTime: 0, + lastSelected: 0, + }, +}; + const MOCK_SNAP_ACCOUNT_2: InternalAccount = { id: 'mock-snap-id-2', address: '0x789', @@ -192,7 +226,7 @@ const MOCK_SNAP_ACCOUNT_2: InternalAccount = { }, }; -const MOCK_TRX_ACCOUNT_1: InternalAccount = { +const MOCK_TRX_ACCOUNT_1: Bip44Account = { id: 'mock-trx-id-1', address: 'TROn11', options: { @@ -230,8 +264,198 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = { }, }; +const MOCK_SIMPLE_ACCOUNT_1: InternalAccount = { + id: 'mock-simple-account-id-1', + address: '0xabc', + options: {}, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Imported Account 1', + keyring: { type: KeyringTypes.simple }, + importTime: 0, + lastSelected: 0, + }, +}; + +const MOCK_SIMPLE_ACCOUNT_2: InternalAccount = { + id: 'mock-simple-account-id-2', + address: '0xdef', + options: {}, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Imported Account 2', + keyring: { type: KeyringTypes.simple }, + importTime: 0, + lastSelected: 0, + }, +}; + +class MockMultichainAccountGroup + implements Partial>> +{ + readonly type = AccountGroupType.MultichainAccount as const; + + readonly id: MultichainAccountGroupId; + + readonly groupIndex: number; + + readonly wallet: MultichainAccountWallet>; + + // We do not make it private so we can push new accounts during tests too (which + // is not how those groups work normally, but that's ok for testing purposes). + readonly accounts: Bip44Account[]; + + constructor({ + wallet, + groupIndex, + accounts, + }: { + wallet: MockMultichainAccountWallet; + groupIndex: number; + accounts: Bip44Account[]; + }) { + this.groupIndex = groupIndex; + this.id = toMultichainAccountGroupId(wallet.id, groupIndex); + this.accounts = accounts; + // This is expected, we only mock partially this wallet, so we need to upcast it. + // NOTE: We cannot use `MockMultichainAccountWallet` here either since it would + // conflict with the expected type from `MultichainAccountGroup.wallet` base type. + this.wallet = wallet as unknown as MultichainAccountWallet< + Bip44Account + >; + } + + getAccounts(): Bip44Account[] { + return this.accounts; + } +} + +class MockMultichainAccountWallet + implements Partial>> +{ + readonly type = AccountWalletType.Entropy as const; + + readonly status = 'ready'; + + readonly id: MultichainAccountWalletId; + + readonly entropySource: EntropySourceId; + + // We do not make it private so we can push new accounts during tests too (which + // is not how those wallets work normally, but that's ok for testing purposes). + readonly groups: Record; + + constructor({ + entropySource, + groups, + }: { + entropySource: EntropySourceId; + groups: Record[]>; + }) { + this.entropySource = entropySource; + this.id = toMultichainAccountWalletId(entropySource); + this.groups = Object.fromEntries( + Object.entries(groups).map(([groupIndex, accounts]) => { + return [ + Number(groupIndex), + new MockMultichainAccountGroup({ + wallet: this, + groupIndex: Number(groupIndex), + accounts, + }), + ]; + }), + ); + } + + getMultichainAccountGroups(): MultichainAccountGroup< + Bip44Account + >[] { + // We're mocking partially this, so we have to cast. + return Object.values(this.groups) as unknown as MultichainAccountGroup< + Bip44Account + >[]; + } +} + +// FIXME: The type used by the service actually conflict with the `account-api`, this has to be +// fixed at the service level. +type MultichainAccountServiceWallet = ReturnType< + MultichainAccountService['getMultichainAccountWallet'] +>; +type MultichainAccountServiceGroup = MultichainAccountGroup< + Bip44Account +>; + const mockGetSelectedMultichainAccountActionHandler = jest.fn(); +/** + * Cast mock multichain account wallet object to a proper multichain account wallet type. + * + * @param wallet - Mock multichain account wallet. + * @returns Multichain account wallet with proper type. + */ +function asMultichainAccountWallet( + wallet: MockMultichainAccountWallet, +): MultichainAccountServiceWallet { + // We're mocking partially this, so we have to cast. + return wallet as unknown as MultichainAccountServiceWallet; +} + +/** + * Cast mock multichain account group object to a proper multichain account group type. + * + * @param group - Mock multichain account group. + * @returns Multichain account group with proper type. + */ +function asMultichainAccountGroup( + group: MockMultichainAccountGroup, +): MultichainAccountServiceGroup { + // We're mocking partially this, so we have to cast. + return group as unknown as MultichainAccountServiceGroup; +} + +/** + * Get multichain account wallets from an account list. + * + * @param accounts - Account list. + * @returns Multichain account wallets. + */ +function getMockMultichainAccountWallets( + accounts: InternalAccount[], +): MultichainAccountServiceWallet[] { + const wallets = accounts.reduce( + (acc, account) => { + if (isBip44Account(account)) { + const { id: entropySource, groupIndex } = account.options.entropy; + + acc[entropySource] ??= {}; + acc[entropySource][groupIndex] ??= []; + acc[entropySource][groupIndex].push(account); + } + + return acc; + }, + {} as Record< + EntropySourceId, + Record[]> + >, + ); + + return Object.entries(wallets).map(([entropySource, groups]) => + asMultichainAccountWallet( + new MockMultichainAccountWallet({ + entropySource, + groups, + }), + ), + ); +} + /** * Sets up the AccountTreeController for testing. * @@ -304,6 +528,10 @@ function setup({ getSelectedMultichainAccount: jest.Mock; getAccount: jest.Mock; }; + MultichainAccountService: { + wallets: MultichainAccountServiceWallet[]; + getMultichainAccountWallets: jest.Mock; + }; UserStorageController: { performGetStorage: jest.Mock; performGetStorageAllFeatureEntries: jest.Mock; @@ -327,6 +555,10 @@ function setup({ getAccount: jest.fn(), getSelectedMultichainAccount: jest.fn(), }, + MultichainAccountService: { + wallets: getMockMultichainAccountWallets(accounts), + getMultichainAccountWallets: jest.fn(), + }, UserStorageController: { getState: jest.fn(), performGetStorage: jest.fn(), @@ -410,6 +642,14 @@ function setup({ 'UserStorageController:performBatchSetStorage', mocks.UserStorageController.performBatchSetStorage, ); + + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => mocks.MultichainAccountService.wallets, + ); + messenger.registerActionHandler( + 'MultichainAccountService:getMultichainAccountWallets', + mocks.MultichainAccountService.getMultichainAccountWallets, + ); } if (keyrings) { @@ -455,7 +695,7 @@ describe('AccountTreeController', () => { accounts: [ MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2, - MOCK_SNAP_ACCOUNT_1, // Belongs to MOCK_HD_ACCOUNT_2's wallet due to shared entropySource + MOCK_SOL_ACCOUNT_1, // Belongs to MOCK_HD_ACCOUNT_2's wallet due to shared entropySource MOCK_SNAP_ACCOUNT_2, // Has its own Snap wallet MOCK_HARDWARE_ACCOUNT_1, // Has its own Keyring wallet ], @@ -489,7 +729,7 @@ describe('AccountTreeController', () => { ); const expectedWalletId2Group2 = toMultichainAccountGroupId( expectedWalletId2, - MOCK_SNAP_ACCOUNT_1.options.entropy.groupIndex, + MOCK_SOL_ACCOUNT_1.options.entropy.groupIndex, ); const expectedSnapWalletId = toAccountWalletId( AccountWalletType.Snap, @@ -558,12 +798,11 @@ describe('AccountTreeController', () => { [expectedWalletId2Group2]: { id: expectedWalletId2Group2, type: AccountGroupType.MultichainAccount, - accounts: [MOCK_SNAP_ACCOUNT_1.id], + accounts: [MOCK_SOL_ACCOUNT_1.id], metadata: { name: 'Account 2', // Updated: per-wallet sequential numbering (wallet 2, account 2) entropy: { - groupIndex: - MOCK_SNAP_ACCOUNT_1.options.entropy.groupIndex, + groupIndex: MOCK_SOL_ACCOUNT_1.options.entropy.groupIndex, }, pinned: false, hidden: false, @@ -772,7 +1011,6 @@ describe('AccountTreeController', () => { MOCK_SNAP_1.id, ); - // FIXME: Do we really want this behavior? expect( controller.state.accountTree.wallets[wallet1Id]?.metadata.name, ).toBe('mock-snap-id-1'); @@ -837,11 +1075,14 @@ describe('AccountTreeController', () => { defaultAccountGroupId, ); - mocks.AccountsController.accounts = [MOCK_HD_ACCOUNT_2]; mocks.KeyringController.keyrings = [MOCK_HD_KEYRING_2]; + mocks.AccountsController.accounts = [MOCK_HD_ACCOUNT_2]; mocks.AccountsController.getSelectedMultichainAccount.mockImplementation( () => MOCK_HD_ACCOUNT_2, ); + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); controller.reinit(); @@ -902,11 +1143,17 @@ describe('AccountTreeController', () => { // We now change the list of accounts entirely and call re-init to re-fetch // the new account list. mocks.AccountsController.accounts = [MOCK_HD_ACCOUNT_2]; + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); controller.reinit(); expect( mocks.AccountsController.listMultichainAccounts, ).toHaveBeenCalledTimes(2); + expect( + mocks.MultichainAccountService.getMultichainAccountWallets, + ).toHaveBeenCalledTimes(2); expect( mocks.AccountsController.getSelectedMultichainAccount, ).toHaveBeenCalledTimes(2); @@ -943,12 +1190,15 @@ describe('AccountTreeController', () => { }); const now = Date.now(); - mocks.AccountsController.listMultichainAccounts.mockReturnValue([ + mocks.AccountsController.accounts = [ // Faking accounts to be out of order: mockAccountWith(1, now + 1000), mockAccountWith(2, now + 2000), mockAccountWith(0, now), - ]); + ]; + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); controller.init(); @@ -1032,10 +1282,10 @@ describe('AccountTreeController', () => { it('selects account with a selector', () => { const mockSolAccount1: Bip44Account = { - ...MOCK_SNAP_ACCOUNT_1, + ...MOCK_SOL_ACCOUNT_1, options: { entropy: { - ...MOCK_SNAP_ACCOUNT_1.options.entropy, + ...MOCK_SOL_ACCOUNT_1.options.entropy, groupIndex: 0, }, }, @@ -1075,73 +1325,50 @@ describe('AccountTreeController', () => { describe('on AccountsController:accountRemoved', () => { it('removes an account from the tree', () => { - // 2 accounts that share the same entropy source (thus, same wallet). - const mockHdAccount1: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, - }; - const mockHdAccount2 = { - ...MOCK_HD_ACCOUNT_2, - options: { - ...MOCK_HD_ACCOUNT_2.options, - entropy: { - ...MOCK_HD_ACCOUNT_2.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, - }; - const { controller, messenger } = setup({ - accounts: [mockHdAccount1, mockHdAccount2], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); // Create entropy wallets that will both get "Wallet" as base name, then get numbered controller.init(); - messenger.publish('AccountsController:accountRemoved', mockHdAccount1.id); + messenger.publish( + 'AccountsController:accountRemoved', + MOCK_SIMPLE_ACCOUNT_2.id, + ); - const walletId1 = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, + const walletId1 = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_ACCOUNT_1.metadata.keyring.type, ); - const walletId1Group = toMultichainAccountGroupId( + const walletId1Group = toAccountGroupId( walletId1, - mockHdAccount1.options.entropy.groupIndex, + MOCK_SIMPLE_ACCOUNT_1.address, ); expect(controller.state).toStrictEqual({ accountTree: { wallets: { [walletId1]: { id: walletId1, - type: AccountWalletType.Entropy, + type: AccountWalletType.Keyring, status: 'ready', groups: { [walletId1Group]: { id: walletId1Group, - type: AccountGroupType.MultichainAccount, + type: AccountGroupType.SingleAccount, metadata: { - name: 'Account 1', - entropy: { - groupIndex: mockHdAccount1.options.entropy.groupIndex, - }, + name: MOCK_SIMPLE_ACCOUNT_1.metadata.name, pinned: false, hidden: false, }, - accounts: [mockHdAccount2.id], // HD account 1 got removed. + accounts: [MOCK_SIMPLE_ACCOUNT_1.id], // HD account 1 got removed. }, }, metadata: { - name: 'Wallet 1', - entropy: { - id: MOCK_HD_KEYRING_1.metadata.id, + name: 'Imported accounts', + keyring: { + type: MOCK_SIMPLE_ACCOUNT_1.metadata.keyring.type, }, }, }, @@ -1154,7 +1381,7 @@ describe('AccountTreeController', () => { // Account groups now get metadata entries during init [walletId1Group]: { name: { - value: 'Account 1', + value: MOCK_SIMPLE_ACCOUNT_1.metadata.name, lastUpdatedAt: expect.any(Number), }, pinned: { @@ -1172,34 +1399,26 @@ describe('AccountTreeController', () => { }); it('prunes an empty group if it holds no accounts', () => { - const mockHdAccount1: Bip44Account = MOCK_HD_ACCOUNT_1; - const mockHdAccount2 = { - ...MOCK_HD_ACCOUNT_2, - options: { - entropy: { - ...MOCK_HD_ACCOUNT_2.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 1, - }, - }, - }; - const { controller, messenger } = setup({ - accounts: [mockHdAccount1, mockHdAccount2], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); - messenger.publish('AccountsController:accountRemoved', mockHdAccount1.id); + messenger.publish( + 'AccountsController:accountRemoved', + MOCK_SIMPLE_ACCOUNT_1.id, + ); - const walletId1 = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, + const walletId1 = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_ACCOUNT_2.metadata.keyring.type, ); - const walletId1Group2 = toMultichainAccountGroupId( + const walletId1Group2 = toAccountGroupId( walletId1, - mockHdAccount2.options.entropy.groupIndex, + MOCK_SIMPLE_ACCOUNT_2.address, ); expect(controller.state).toStrictEqual({ @@ -1207,28 +1426,25 @@ describe('AccountTreeController', () => { wallets: { [walletId1]: { id: walletId1, - type: AccountWalletType.Entropy, + type: AccountWalletType.Keyring, status: 'ready', groups: { // First group gets removed as a result of pruning. [walletId1Group2]: { id: walletId1Group2, - type: AccountGroupType.MultichainAccount, + type: AccountGroupType.SingleAccount, metadata: { - name: 'Account 2', - entropy: { - groupIndex: mockHdAccount2.options.entropy.groupIndex, - }, + name: MOCK_SIMPLE_ACCOUNT_2.metadata.name, pinned: false, hidden: false, }, - accounts: [mockHdAccount2.id], + accounts: [MOCK_SIMPLE_ACCOUNT_2.id], }, }, metadata: { - name: 'Wallet 1', - entropy: { - id: MOCK_HD_KEYRING_1.metadata.id, + name: 'Imported accounts', + keyring: { + type: MOCK_SIMPLE_ACCOUNT_2.metadata.keyring.type, }, }, }, @@ -1241,7 +1457,7 @@ describe('AccountTreeController', () => { // Both groups get metadata during init, but first group metadata gets cleaned up when pruned [walletId1Group2]: { name: { - value: 'Account 2', // This is the second account in the wallet + value: MOCK_SIMPLE_ACCOUNT_2.metadata.name, lastUpdatedAt: expect.any(Number), }, pinned: { @@ -1259,16 +1475,17 @@ describe('AccountTreeController', () => { }); it('prunes an empty wallet if it holds no groups', () => { - const mockHdAccount1: Bip44Account = MOCK_HD_ACCOUNT_1; - const { controller, messenger } = setup({ - accounts: [mockHdAccount1], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); - messenger.publish('AccountsController:accountRemoved', mockHdAccount1.id); + messenger.publish( + 'AccountsController:accountRemoved', + MOCK_SIMPLE_ACCOUNT_1.id, + ); expect(controller.state).toStrictEqual({ accountGroupsMetadata: {}, @@ -1284,17 +1501,16 @@ describe('AccountTreeController', () => { }); it('prunes custom wallet metadata when wallet is removed', () => { - const mockHdAccount1: Bip44Account = MOCK_HD_ACCOUNT_1; - const { controller, messenger } = setup({ - accounts: [mockHdAccount1], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); - const walletId = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, + const walletId = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, ); // Set custom wallet name @@ -1309,7 +1525,10 @@ describe('AccountTreeController', () => { }); // Remove the account, which should prune the wallet and its metadata - messenger.publish('AccountsController:accountRemoved', mockHdAccount1.id); + messenger.publish( + 'AccountsController:accountRemoved', + MOCK_SIMPLE_ACCOUNT_1.id, + ); // Verify both wallet and its metadata are completely removed expect(controller.state.accountTree.wallets[walletId]).toBeUndefined(); @@ -1319,7 +1538,8 @@ describe('AccountTreeController', () => { it('does not remove account if init has not been called', () => { const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); // Force ref to the controller, even if we don't use it in this test. @@ -1333,7 +1553,7 @@ describe('AccountTreeController', () => { messenger.publish( 'AccountsController:accountRemoved', - MOCK_HD_ACCOUNT_1.id, + MOCK_SIMPLE_ACCOUNT_1.id, ); expect(mockAccountTreeChange).not.toHaveBeenCalled(); @@ -1345,12 +1565,12 @@ describe('AccountTreeController', () => { const evmAccount = MOCK_HD_ACCOUNT_1; const solAccount = { - ...MOCK_SNAP_ACCOUNT_1, + ...MOCK_SOL_ACCOUNT_1, id: 'mock-sol-id-1', options: { - ...MOCK_SNAP_ACCOUNT_1.options, + ...MOCK_SOL_ACCOUNT_1.options, entropy: { - ...MOCK_SNAP_ACCOUNT_1.options.entropy, + ...MOCK_SOL_ACCOUNT_1.options.entropy, id: MOCK_HD_KEYRING_1.metadata.id, groupIndex: 0, derivationPath: '', @@ -1360,23 +1580,48 @@ describe('AccountTreeController', () => { const tronAccount = MOCK_TRX_ACCOUNT_1; - const { controller, messenger } = setup({ + const { controller, messenger, mocks } = setup({ accounts: [], keyrings: [MOCK_HD_KEYRING_1], }); controller.init(); - // Publish in shuffled order: SOL, TRON, EVM - messenger.publish('AccountsController:accountAdded', solAccount); - messenger.publish('AccountsController:accountAdded', tronAccount); - messenger.publish('AccountsController:accountAdded', evmAccount); - - const walletId = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, - ); + const entropySource = MOCK_HD_KEYRING_1.metadata.id; + const walletId = toMultichainAccountWalletId(entropySource); const groupId = toMultichainAccountGroupId(walletId, 0); + const mockWallet = new MockMultichainAccountWallet({ + entropySource, + groups: { + 0: [], // No accounts for now. + }, + }); + + const mockGroup = mockWallet.groups[0]; + expect(mockGroup).toBeDefined(); + expect(mockGroup.id).toBe(groupId); + + mocks.MultichainAccountService.wallets = [ + asMultichainAccountWallet(mockWallet), + ]; + + const addAccountToGroup = (account: Bip44Account) => { + mocks.AccountsController.accounts.push(account); + + mockGroup.accounts.push(account); + + messenger.publish( + 'MultichainAccountService:multichainAccountGroupUpdated', + asMultichainAccountGroup(mockGroup), + ); + }; + + // Publish in shuffled order: SOL, TRON, EVM + addAccountToGroup(solAccount); + addAccountToGroup(tronAccount); + addAccountToGroup(evmAccount); + const group = controller.state.accountTree.wallets[walletId]?.groups[groupId]; expect(group).toBeDefined(); @@ -1392,46 +1637,29 @@ describe('AccountTreeController', () => { describe('on AccountsController:accountAdded', () => { it('adds an account to the tree', () => { - // 2 accounts that share the same entropy source (thus, same wallet). - const mockHdAccount1: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, - }; - const mockHdAccount2 = { - ...MOCK_HD_ACCOUNT_2, - options: { - ...MOCK_HD_ACCOUNT_2.options, - entropy: { - ...MOCK_HD_ACCOUNT_2.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, - }; - const { controller, messenger } = setup({ - accounts: [mockHdAccount1], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); - // Create entropy wallets that will both get "Wallet" as base name, then get numbered controller.init(); - messenger.publish('AccountsController:accountAdded', mockHdAccount2); + messenger.publish( + 'AccountsController:accountAdded', + MOCK_SIMPLE_ACCOUNT_2, + ); - const walletId1 = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, + const walletId1 = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, ); - const walletId1Group = toMultichainAccountGroupId( + const walletId1Group = toAccountGroupId( walletId1, - mockHdAccount1.options.entropy.groupIndex, + MOCK_SIMPLE_ACCOUNT_1.address, + ); + const walletId2Group = toAccountGroupId( + walletId1, + MOCK_SIMPLE_ACCOUNT_2.address, ); expect(controller.state).toStrictEqual({ accountTree: { @@ -1439,37 +1667,60 @@ describe('AccountTreeController', () => { wallets: { [walletId1]: { id: walletId1, - type: AccountWalletType.Entropy, + type: AccountWalletType.Keyring, status: 'ready', groups: { [walletId1Group]: { id: walletId1Group, - type: AccountGroupType.MultichainAccount, + type: AccountGroupType.SingleAccount, metadata: { - name: 'Account 1', - entropy: { - groupIndex: mockHdAccount1.options.entropy.groupIndex, - }, + name: MOCK_SIMPLE_ACCOUNT_1.metadata.name, pinned: false, hidden: false, }, - accounts: [mockHdAccount1.id, mockHdAccount2.id], // HD account 2 got added. + accounts: [MOCK_SIMPLE_ACCOUNT_1.id], + }, + // New group that got added. + [walletId2Group]: { + id: walletId2Group, + type: AccountGroupType.SingleAccount, + metadata: { + name: MOCK_SIMPLE_ACCOUNT_2.metadata.name, + pinned: false, + hidden: false, + }, + accounts: [MOCK_SIMPLE_ACCOUNT_2.id], }, }, metadata: { - name: 'Wallet 1', - entropy: { - id: MOCK_HD_KEYRING_1.metadata.id, + name: 'Imported accounts', + keyring: { + type: MOCK_SIMPLE_ACCOUNT_1.metadata.keyring.type, }, }, }, }, }, accountGroupsMetadata: { - // Account groups now get metadata entries during init + // Account groups now get metadata entries during init. [walletId1Group]: { name: { - value: 'Account 1', + value: MOCK_SIMPLE_ACCOUNT_1.metadata.name, + lastUpdatedAt: expect.any(Number), + }, + pinned: { + value: false, + lastUpdatedAt: 0, + }, + hidden: { + value: false, + lastUpdatedAt: 0, + }, + }, + // New group metadata. + [walletId2Group]: { + name: { + value: MOCK_SIMPLE_ACCOUNT_2.metadata.name, lastUpdatedAt: expect.any(Number), }, pinned: { @@ -1488,57 +1739,100 @@ describe('AccountTreeController', () => { } as AccountTreeControllerState); }); - it('adds a new wallet to the tree', () => { - // 2 accounts that share the same entropy source (thus, same wallet). - const mockHdAccount1: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, - }; - const mockHdAccount2 = { - ...MOCK_HD_ACCOUNT_2, - options: { - ...MOCK_HD_ACCOUNT_2.options, - entropy: { - ...MOCK_HD_ACCOUNT_2.options.entropy, - id: MOCK_HD_KEYRING_2.metadata.id, - groupIndex: 0, - }, - }, - }; + it('skips BIP-44 account', () => { + const { controller, messenger } = setup({ + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], + }); + + controller.init(); + // Adding a BIP-44 accounts should not trigger any change. + const mockAccountTreeChange = jest.fn(); + messenger.subscribe( + 'AccountTreeController:accountTreeChange', + mockAccountTreeChange, + ); + + // BIP-44 accounts are skipped, we only listen to the `MultichainAccountService` change + // for those accounts! + messenger.publish('AccountsController:accountAdded', MOCK_HD_ACCOUNT_1); + + expect(mockAccountTreeChange).not.toHaveBeenCalled(); + }); + }); + + describe('on MultichainAccountService:multichainAccountGroupCreated', () => { + it('prevent event if controller is not init yet', () => { const { controller, messenger, mocks } = setup({ - accounts: [mockHdAccount1], + accounts: [MOCK_HD_ACCOUNT_1], keyrings: [MOCK_HD_KEYRING_1], }); - // Create entropy wallets that will both get "Wallet" as base name, then get numbered - controller.init(); + // We omit `init` on purpose for this test. - mocks.KeyringController.keyrings = [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2]; - mocks.AccountsController.accounts = [mockHdAccount1, mockHdAccount2]; - messenger.publish('AccountsController:accountAdded', mockHdAccount2); + expect(controller.state.accountTree.wallets).toStrictEqual({}); + + const walletId = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, + ); + expect(mocks.MultichainAccountService.wallets).toHaveLength(1); + const wallet = mocks.MultichainAccountService.wallets[0]; + expect(wallet.id).toBe(walletId); + const [group] = wallet.getMultichainAccountGroups(); + expect(group).toBeDefined(); + messenger.publish( + 'MultichainAccountService:multichainAccountGroupCreated', + group, + ); + + // Since the controller has not been initialized yet, the event had been omitted! + expect(controller.state.accountTree.wallets).toStrictEqual({}); + }); + + it('adds a new wallet to the tree', () => { + const { controller, messenger, mocks } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + controller.init(); const walletId1 = toMultichainAccountWalletId( MOCK_HD_KEYRING_1.metadata.id, ); const walletId1Group = toMultichainAccountGroupId( walletId1, - mockHdAccount1.options.entropy.groupIndex, + MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, ); const walletId2 = toMultichainAccountWalletId( MOCK_HD_KEYRING_2.metadata.id, ); const walletId2Group = toMultichainAccountGroupId( walletId2, - mockHdAccount2.options.entropy.groupIndex, + MOCK_HD_ACCOUNT_2.options.entropy.groupIndex, ); + + mocks.KeyringController.keyrings = [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2]; + mocks.AccountsController.accounts = [ + MOCK_HD_ACCOUNT_1, + MOCK_HD_ACCOUNT_2, + ]; + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); + + expect(mocks.MultichainAccountService.wallets).toHaveLength(2); + const wallet2 = mocks.MultichainAccountService.wallets[1]; + expect(wallet2.id).toBe(walletId2); + const [wallet2Group] = wallet2.getMultichainAccountGroups(); + expect(wallet2Group).toBeDefined(); + messenger.publish( + 'MultichainAccountService:multichainAccountGroupCreated', + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + wallet2Group!, + ); + expect(controller.state).toStrictEqual({ accountTree: { wallets: { @@ -1553,12 +1847,12 @@ describe('AccountTreeController', () => { metadata: { name: 'Account 1', entropy: { - groupIndex: mockHdAccount1.options.entropy.groupIndex, + groupIndex: MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, }, pinned: false, hidden: false, }, - accounts: [mockHdAccount1.id], + accounts: [MOCK_HD_ACCOUNT_1.id], }, }, metadata: { @@ -1568,8 +1862,8 @@ describe('AccountTreeController', () => { }, }, }, + // New wallet being added. [walletId2]: { - // New wallet automatically added. id: walletId2, type: AccountWalletType.Entropy, status: 'ready', @@ -1578,14 +1872,14 @@ describe('AccountTreeController', () => { id: walletId2Group, type: AccountGroupType.MultichainAccount, metadata: { - name: 'Account 1', // Updated: per-wallet naming (different wallet) + name: 'Account 1', entropy: { - groupIndex: mockHdAccount2.options.entropy.groupIndex, + groupIndex: MOCK_HD_ACCOUNT_2.options.entropy.groupIndex, }, pinned: false, hidden: false, }, - accounts: [mockHdAccount2.id], + accounts: [MOCK_HD_ACCOUNT_2.id], }, }, metadata: { @@ -1784,11 +2078,11 @@ describe('AccountTreeController', () => { it('updates AccountsController selected account (with non-EVM account) when selectedAccountGroup changes', () => { const nonEvmAccount2 = { - ...MOCK_SNAP_ACCOUNT_1, + ...MOCK_SOL_ACCOUNT_1, options: { - ...MOCK_SNAP_ACCOUNT_1.options, + ...MOCK_SOL_ACCOUNT_1.options, entropy: { - ...MOCK_SNAP_ACCOUNT_1.options.entropy, + ...MOCK_SOL_ACCOUNT_1.options.entropy, id: MOCK_HD_KEYRING_2.metadata.id, // Wallet 2. groupIndex: 0, // Account 1 }, @@ -2039,44 +2333,35 @@ describe('AccountTreeController', () => { it('updates selectedAccountGroup when last account in selected group is removed and other groups exist', () => { const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2], - keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); // Select the first group - const expectedWalletId1 = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, - ); - const expectedGroupId1 = toMultichainAccountGroupId( - expectedWalletId1, - MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, - ); - controller.setSelectedAccountGroup(expectedGroupId1); - - const expectedWalletId2 = toMultichainAccountWalletId( - MOCK_HD_KEYRING_2.metadata.id, - ); - const expectedGroupId2 = toMultichainAccountGroupId( - expectedWalletId2, - MOCK_HD_ACCOUNT_2.options.entropy.groupIndex, + const wallet = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, ); + const group1 = toAccountGroupId(wallet, MOCK_SIMPLE_ACCOUNT_1.address); + const group2 = toAccountGroupId(wallet, MOCK_SIMPLE_ACCOUNT_2.address); + controller.setSelectedAccountGroup(group1); // Remove the account from the selected group - tests true branch and findFirstNonEmptyGroup finding a group messenger.publish( 'AccountsController:accountRemoved', - MOCK_HD_ACCOUNT_1.id, + MOCK_SIMPLE_ACCOUNT_1.id, ); // Should automatically switch to the remaining group (tests findFirstNonEmptyGroup returning a group) - expect(controller.getSelectedAccountGroup()).toBe(expectedGroupId2); + expect(controller.getSelectedAccountGroup()).toBe(group2); }); it('sets selectedAccountGroup to empty when no non-empty groups exist', () => { const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); @@ -2084,7 +2369,7 @@ describe('AccountTreeController', () => { // Remove the only account - tests findFirstNonEmptyGroup returning empty string messenger.publish( 'AccountsController:accountRemoved', - MOCK_HD_ACCOUNT_1.id, + MOCK_SIMPLE_ACCOUNT_1.id, ); // Should fall back to empty string when no groups have accounts @@ -2685,67 +2970,46 @@ describe('AccountTreeController', () => { describe('Fallback Naming', () => { it('uses consistent default naming regardless of account import time', () => { - const mockAccount1: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 0, - }, - }, + const mockAccount1: InternalAccount = { + ...MOCK_SIMPLE_ACCOUNT_1, metadata: { - ...MOCK_HD_ACCOUNT_1.metadata, + ...MOCK_SIMPLE_ACCOUNT_1.metadata, importTime: Date.now() + 1000, + name: '', // Empty name to trigger naming logic }, }; - const mockAccount2: Bip44Account = { - ...MOCK_HD_ACCOUNT_2, - options: { - ...MOCK_HD_ACCOUNT_2.options, - entropy: { - ...MOCK_HD_ACCOUNT_2.options.entropy, - id: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: 1, - }, - }, + const mockAccount2: InternalAccount = { + ...MOCK_SIMPLE_ACCOUNT_2, metadata: { - ...MOCK_HD_ACCOUNT_2.metadata, + ...MOCK_SIMPLE_ACCOUNT_2.metadata, importTime: Date.now() - 1000, + name: '', // Empty name to trigger naming logic }, }; const { controller } = setup({ accounts: [mockAccount2, mockAccount1], - keyrings: [MOCK_HD_KEYRING_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); controller.init(); - const expectedWalletId = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, - ); - - const expectedGroupId1 = toMultichainAccountGroupId( - expectedWalletId, - mockAccount1.options.entropy.groupIndex, - ); - - const expectedGroupId2 = toMultichainAccountGroupId( - expectedWalletId, - mockAccount2.options.entropy.groupIndex, + const walletId = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, ); + const groupId1 = toAccountGroupId(walletId, mockAccount1.address); + const groupId2 = toAccountGroupId(walletId, mockAccount2.address); - const wallet = controller.state.accountTree.wallets[expectedWalletId]; - const group1 = wallet?.groups[expectedGroupId1]; - const group2 = wallet?.groups[expectedGroupId2]; + const wallet = controller.state.accountTree.wallets[walletId]; + const group1 = wallet?.groups[groupId1]; + const group2 = wallet?.groups[groupId2]; // Groups should use consistent default naming regardless of import time // Updated expectations based on per-wallet sequential naming logic - expect(group1?.metadata.name).toBe('Account 2'); // Updated: reflects actual naming logic - expect(group2?.metadata.name).toBe('Account 1'); // Updated: reflects actual naming logic + expect(group1?.metadata.name).toBe('Imported Account 2'); // Updated: reflects actual naming logic + expect(group2?.metadata.name).toBe('Imported Account 1'); // Updated: reflects actual naming logic }); it('uses fallback naming when rule-based naming returns empty string', () => { @@ -2853,18 +3117,32 @@ describe('AccountTreeController', () => { // Add the new account to the existing group mocks.AccountsController.accounts = [existingAccount, newAccount]; - messenger.publish('AccountsController:accountAdded', newAccount); + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); - const expectedWalletId = toMultichainAccountWalletId( + const walletId = toMultichainAccountWalletId( MOCK_HD_KEYRING_1.metadata.id, ); - const expectedGroupId = toMultichainAccountGroupId( - expectedWalletId, + const groupId = toMultichainAccountGroupId( + walletId, 0, // Same group index ); - const wallet = controller.state.accountTree.wallets[expectedWalletId]; - const group = wallet?.groups[expectedGroupId]; + expect(mocks.MultichainAccountService.wallets).toHaveLength(1); + const existingWallet = mocks.MultichainAccountService.wallets[0]; + expect(existingWallet.id).toBe(walletId); + const [existingGroup] = existingWallet.getMultichainAccountGroups(); + expect(existingGroup).toBeDefined(); + expect(existingGroup.getAccounts()).toHaveLength(2); + messenger.publish( + 'MultichainAccountService:multichainAccountGroupUpdated', + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + existingGroup!, + ); + + const wallet = controller.state.accountTree.wallets[walletId]; + const group = wallet?.groups[groupId]; // The group should use consistent default naming expect(group?.metadata.name).toBe('Account 1'); @@ -3120,59 +3398,32 @@ describe('AccountTreeController', () => { // This test verifies that account names remain consistent after restart // and don't change from "Account 1" to "Account 2" etc. - // Create two accounts in the same wallet but different groups - const account1: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - id: 'account-1', - metadata: { - ...MOCK_HD_ACCOUNT_1.metadata, - name: '', // Empty name to force default naming - }, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - groupIndex: 0, - }, - }, - }; - - const account2: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, - id: 'account-2', - address: '0x456', - metadata: { - ...MOCK_HD_ACCOUNT_1.metadata, - name: '', // Empty name to force default naming - }, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - ...MOCK_HD_ACCOUNT_1.options.entropy, - groupIndex: 1, - }, - }, - }; - const { controller, messenger } = setup({ - accounts: [account1, account2], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); // First init - accounts get named controller.init(); - const walletId = toMultichainAccountWalletId( - MOCK_HD_KEYRING_1.metadata.id, + const walletId = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, + ); + const group1Id = toAccountGroupId( + walletId, + MOCK_SIMPLE_ACCOUNT_1.address, + ); + const group2Id = toAccountGroupId( + walletId, + MOCK_SIMPLE_ACCOUNT_2.address, ); - const group1Id = toMultichainAccountGroupId(walletId, 0); - const group2Id = toMultichainAccountGroupId(walletId, 1); // Check initial names (both groups use entropy.groupIndex) const state1 = controller.state; const wallet1 = state1.accountTree.wallets[walletId]; - expect(wallet1.groups[group1Id].metadata.name).toBe('Account 1'); // groupIndex 0 → Account 1 - expect(wallet1.groups[group2Id].metadata.name).toBe('Account 2'); // groupIndex 1 → Account 2 + expect(wallet1.groups[group1Id].metadata.name).toBe('Imported Account 1'); + expect(wallet1.groups[group2Id].metadata.name).toBe('Imported Account 2'); // Simulate app restart by re-initializing controller.reinit(); @@ -3180,36 +3431,27 @@ describe('AccountTreeController', () => { // Names should remain the same (consistent entropy.groupIndex) const state2 = controller.state; const wallet2 = state2.accountTree.wallets[walletId]; - expect(wallet2.groups[group1Id].metadata.name).toBe('Account 1'); - expect(wallet2.groups[group2Id].metadata.name).toBe('Account 2'); + expect(wallet2.groups[group1Id].metadata.name).toBe('Imported Account 1'); + expect(wallet2.groups[group2Id].metadata.name).toBe('Imported Account 2'); // Add a new account after restart - const newAccount: Bip44Account = { - ...MOCK_HD_ACCOUNT_1, + const newAccount: InternalAccount = { + ...MOCK_SIMPLE_ACCOUNT_1, id: 'new-account', address: '0xNEW', metadata: { - ...MOCK_HD_ACCOUNT_1.metadata, + ...MOCK_SIMPLE_ACCOUNT_1.metadata, name: '', // Empty to force default naming }, - options: { - ...MOCK_HD_ACCOUNT_1.options, - entropy: { - type: 'mnemonic', - id: MOCK_HD_KEYRING_1.metadata.id, - derivationPath: "m/44'/60'/2'/0/0", - groupIndex: 2, - }, - }, }; messenger.publish('AccountsController:accountAdded', newAccount); // New account should get Account 3, not duplicate an existing name - const group3Id = toMultichainAccountGroupId(walletId, 2); + const group3Id = toAccountGroupId(walletId, newAccount.address); const state3 = controller.state; const wallet3 = state3.accountTree.wallets[walletId]; - expect(wallet3.groups[group3Id].metadata.name).toBe('Account 3'); + expect(wallet3.groups[group3Id].metadata.name).toBe('Imported Account 3'); // All names should be different const allNames = [ @@ -3781,8 +4023,8 @@ describe('AccountTreeController', () => { it('emits accountTreeChange when account is added', () => { const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1], + keyrings: [MOCK_SIMPLE_KEYRING], }); const accountTreeChangeListener = jest.fn(); @@ -3795,7 +4037,7 @@ describe('AccountTreeController', () => { jest.clearAllMocks(); messenger.publish('AccountsController:accountAdded', { - ...MOCK_HD_ACCOUNT_2, + ...MOCK_SIMPLE_ACCOUNT_2, }); expect(accountTreeChangeListener).toHaveBeenCalledWith( @@ -3806,8 +4048,8 @@ describe('AccountTreeController', () => { it('emits accountTreeChange when account is removed', () => { const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2], - keyrings: [MOCK_HD_KEYRING_1], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); const accountTreeChangeListener = jest.fn(); @@ -3821,7 +4063,7 @@ describe('AccountTreeController', () => { messenger.publish( 'AccountsController:accountRemoved', - MOCK_HD_ACCOUNT_2.id, + MOCK_SIMPLE_ACCOUNT_2.id, ); expect(accountTreeChangeListener).toHaveBeenCalledWith( @@ -3833,8 +4075,8 @@ describe('AccountTreeController', () => { it('emits selectedAccountGroupChange when account removal causes empty group and auto-selection', () => { // Set up with two accounts in different groups to ensure group change on removal const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1, MOCK_SNAP_ACCOUNT_1], - keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2], + accounts: [MOCK_SIMPLE_ACCOUNT_1, MOCK_SIMPLE_ACCOUNT_2], + keyrings: [MOCK_SIMPLE_KEYRING], }); const selectedAccountGroupChangeListener = jest.fn(); @@ -3846,10 +4088,11 @@ describe('AccountTreeController', () => { controller.init(); // Set selected group to be the group we're about to empty - const walletId = toMultichainAccountWalletId( - MOCK_HD_KEYRING_2.metadata.id, + const walletId = toAccountWalletId( + AccountWalletType.Keyring, + MOCK_SIMPLE_KEYRING.type, ); - const groupId = toMultichainAccountGroupId(walletId, 1); + const groupId = toAccountGroupId(walletId, MOCK_SIMPLE_ACCOUNT_2.address); controller.setSelectedAccountGroup(groupId); jest.clearAllMocks(); @@ -3857,7 +4100,7 @@ describe('AccountTreeController', () => { // Remove the only account in the selected group, which should trigger auto-selection messenger.publish( 'AccountsController:accountRemoved', - MOCK_SNAP_ACCOUNT_1.id, + MOCK_SIMPLE_ACCOUNT_2.id, ); const newSelectedGroup = @@ -3926,11 +4169,14 @@ describe('AccountTreeController', () => { selectedAccountGroupChangeListener, ); - mocks.AccountsController.accounts = [MOCK_HD_ACCOUNT_2]; mocks.KeyringController.keyrings = [MOCK_HD_KEYRING_2]; + mocks.AccountsController.accounts = [MOCK_HD_ACCOUNT_2]; mocks.AccountsController.getSelectedMultichainAccount.mockImplementation( () => MOCK_HD_ACCOUNT_2, ); + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); controller.reinit(); @@ -3952,7 +4198,7 @@ describe('AccountTreeController', () => { it('emits selectedAccountGroupChange when setSelectedAccountGroup is called', () => { // Use different keyring types to ensure different groups const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1, MOCK_SNAP_ACCOUNT_1], + accounts: [MOCK_HD_ACCOUNT_1, MOCK_SOL_ACCOUNT_1], keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2], }); @@ -3985,7 +4231,7 @@ describe('AccountTreeController', () => { it('emits selectedAccountGroupChange when selected account changes via AccountsController', () => { // Use different keyring types to ensure different groups const { controller, messenger } = setup({ - accounts: [MOCK_HD_ACCOUNT_1, MOCK_SNAP_ACCOUNT_1], + accounts: [MOCK_HD_ACCOUNT_1, MOCK_SOL_ACCOUNT_1], keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2], }); @@ -4004,7 +4250,7 @@ describe('AccountTreeController', () => { messenger.publish( 'AccountsController:selectedAccountChange', - MOCK_SNAP_ACCOUNT_1, + MOCK_SOL_ACCOUNT_1, ); const newSelectedGroup = @@ -4635,6 +4881,9 @@ describe('AccountTreeController', () => { // Then, we insert a second account (index 1), but we re-order it so it appears // before the first account (index 0). mocks.AccountsController.accounts = [mockHdAccount2, mockHdAccount1]; + mocks.MultichainAccountService.wallets = getMockMultichainAccountWallets( + mocks.AccountsController.accounts, + ); // Re-init the controller should still give proper naming. controller.reinit(); diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index 2180076e233..8dc62d32c76 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -1,16 +1,24 @@ -import { AccountWalletType, select } from '@metamask/account-api'; +import { + AccountWalletType, + isBip44Account, + select, +} from '@metamask/account-api'; +import { AccountGroupType } from '@metamask/account-api'; import type { + MultichainAccountWalletStatus, AccountGroupId, AccountWalletId, AccountSelector, MultichainAccountWalletId, - AccountGroupType, + Bip44Account, + MultichainAccountWallet, + MultichainAccountGroup, } from '@metamask/account-api'; -import type { MultichainAccountWalletStatus } 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'; @@ -22,7 +30,11 @@ import { } from './backup-and-sync/analytics'; import { BackupAndSyncService } from './backup-and-sync/service'; import type { BackupAndSyncContext } from './backup-and-sync/types'; -import type { AccountGroupObject, AccountTypeOrderKey } from './group'; +import type { + AccountGroupMultichainAccountObject, + AccountGroupObject, + AccountTypeOrderKey, +} from './group'; import { ACCOUNT_TYPE_TO_SORT_ORDER, isAccountGroupNameUnique, @@ -30,8 +42,10 @@ import { MAX_SORT_ORDER, } from './group'; import { projectLogger as log } from './logger'; -import type { Rule } from './rule'; -import { EntropyRule } from './rules/entropy'; +import { + getEntropyDefaultAccountGroupPrefix, + getEntropyDefaultAccountWalletName, +} from './rules/entropy'; import { KeyringRule } from './rules/keyring'; import { SnapRule } from './rules/snap'; import type { @@ -40,6 +54,7 @@ import type { AccountTreeControllerMessenger, AccountTreeControllerState, } from './types'; +import type { AccountWalletEntropyObject } from './wallet'; import { type AccountWalletObject, type AccountWalletObjectOf } from './wallet'; export const controllerName = 'AccountTreeController'; @@ -130,7 +145,7 @@ export class AccountTreeController extends BaseController< */ readonly #backupAndSyncService: BackupAndSyncService; - readonly #rules: [EntropyRule, SnapRule, KeyringRule]; + readonly #rules: [SnapRule, KeyringRule]; readonly #trace: TraceCallback; @@ -185,7 +200,7 @@ export class AccountTreeController extends BaseController< // Rules to apply to construct the wallets tree. this.#rules = [ // 1. We group by entropy-source - new EntropyRule(this.messenger), + // NOTE: This is now done by consuming the `MultichainAccountService` wallets and groups. // 2. We group by Snap ID new SnapRule(this.messenger), // 3. We group by wallet type (this rule cannot fail and will group all non-matching accounts) @@ -240,6 +255,20 @@ export class AccountTreeController extends BaseController< }, ); + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupCreated', + (group) => { + this.#handleMultichainAccountGroupCreatedOrUpdated(group); + }, + ); + + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupUpdated', + (group) => { + this.#handleMultichainAccountGroupCreatedOrUpdated(group); + }, + ); + this.messenger.subscribe( 'MultichainAccountService:walletStatusChange', (walletId, status) => { @@ -266,6 +295,7 @@ export class AccountTreeController extends BaseController< log('Initializing...'); + // For now, we always re-compute all wallets, we do not re-use the existing state. const wallets: AccountTreeControllerState['accountTree']['wallets'] = {}; // Clear mappings for fresh rebuild. @@ -277,6 +307,17 @@ export class AccountTreeController extends BaseController< const previousSelectedAccountGroup = this.state.accountTree.selectedAccountGroup; + // Insert all multichain account wallet/groups. + for (const wallet of this.#getMultichainAccountWallets()) { + for (const group of wallet.getMultichainAccountGroups()) { + this.#insertOrUpdateMultichainAccountWalletAndGroup( + wallets, + wallet, + group, + ); + } + } + // There's no guarantee that accounts would be sorted by their import time // with `listMultichainAccounts`. We have to sort them here before constructing // the tree. @@ -292,9 +333,15 @@ export class AccountTreeController extends BaseController< (a, b) => a.metadata.importTime - b.metadata.importTime, ); - // For now, we always re-compute all wallets, we do not re-use the existing state. + // Insert all other kind of accounts (private keys, HW, other Snap accounts, etc.). for (const account of accounts) { - this.#insert(wallets, account); + // BIP-44 accounts are owned byt the `MultichainAccountService`, so we have to skip them + // here, since we've already added wallets/groups for each of them. + if (isBip44Account(account)) { + continue; + } + + this.#insertAccount(wallets, account); } // Once we have the account tree, we can apply persisted metadata (names + UI states). @@ -378,22 +425,13 @@ export class AccountTreeController extends BaseController< this.init(); } - /** - * Rule for entropy-base wallets. - * - * @returns The rule for entropy-based wallets. - */ - #getEntropyRule(): EntropyRule { - return this.#rules[0]; - } - /** * Rule for Snap-base wallets. * * @returns The rule for snap-based wallets. */ #getSnapRule(): SnapRule { - return this.#rules[1]; + return this.#rules[0]; } /** @@ -406,7 +444,26 @@ export class AccountTreeController extends BaseController< * any other rules. */ #getKeyringRule(): KeyringRule { - return this.#rules[2]; + return this.#rules[1]; + } + + /** + * Gets the default name for a wallet. + * + * @param wallet - The wallet object to get the name for. + * @returns The name for this wallet. + */ + #getDefaultAccountWalletName( + wallet: AccountWalletObjectOf, + ): string { + switch (wallet.type) { + case AccountWalletType.Entropy: + return getEntropyDefaultAccountWalletName(this.messenger, wallet); + case AccountWalletType.Snap: + return this.#getSnapRule().getDefaultAccountWalletName(wallet); + default: + return this.#getKeyringRule().getDefaultAccountWalletName(wallet); + } } /** @@ -429,48 +486,11 @@ export class AccountTreeController extends BaseController< wallet.metadata.name = persistedMetadata.name.value; } else if (!wallet.metadata.name) { // Generate default name if none exists - if (wallet.type === AccountWalletType.Entropy) { - wallet.metadata.name = - this.#getEntropyRule().getDefaultAccountWalletName(wallet); - } else if (wallet.type === AccountWalletType.Snap) { - wallet.metadata.name = - this.#getSnapRule().getDefaultAccountWalletName(wallet); - } else { - wallet.metadata.name = - this.#getKeyringRule().getDefaultAccountWalletName(wallet); - } + wallet.metadata.name = this.#getDefaultAccountWalletName(wallet); log(`[${wallet.id}] Set default name to: "${wallet.metadata.name}"`); } } - /** - * Gets the appropriate rule instance for a given wallet type. - * - * @param wallet - The wallet object to get the rule for. - * @returns The rule instance that handles the wallet's type. - */ - #getRuleForWallet( - wallet: AccountWalletObjectOf, - ): Rule { - switch (wallet.type) { - case AccountWalletType.Entropy: - return this.#getEntropyRule() as unknown as Rule< - WalletType, - AccountGroupType - >; - case AccountWalletType.Snap: - return this.#getSnapRule() as unknown as Rule< - WalletType, - AccountGroupType - >; - default: - return this.#getKeyringRule() as unknown as Rule< - WalletType, - AccountGroupType - >; - } - } - /** * Gets the computed name of a group (using its associated accounts). * @@ -514,6 +534,25 @@ export class AccountTreeController extends BaseController< return proposedName; } + /** + * Gets the default account group prefix for a wallet. + * + * @param wallet - The wallet object to get the prefix for. + * @returns The prefix for this wallet. + */ + #getDefaultAccountGroupPrefix( + wallet: AccountWalletObjectOf, + ): string { + switch (wallet.type) { + case AccountWalletType.Entropy: + return getEntropyDefaultAccountGroupPrefix(); + case AccountWalletType.Snap: + return this.#getSnapRule().getDefaultAccountGroupPrefix(wallet); + default: + return this.#getKeyringRule().getDefaultAccountGroupPrefix(wallet); + } + } + /** * Gets the default name of a group. * @@ -529,11 +568,8 @@ export class AccountTreeController extends BaseController< group: AccountGroupObject, nextNaturalNameIndex?: number, ): string { - // Get the appropriate rule for this wallet type - const rule = this.#getRuleForWallet(wallet); - // Get the prefix for groups of this wallet - const namePrefix = rule.getDefaultAccountGroupPrefix(wallet); + const namePrefix = this.#getDefaultAccountGroupPrefix(wallet); // Parse the highest account index being used (similar to accounts-controller) let highestNameIndex = 0; @@ -808,10 +844,16 @@ export class AccountTreeController extends BaseController< return; } + if (isBip44Account(account)) { + // We're skipping BIP-44 accounts since we rely on the `MultichainAccountService` to do + // the grouping of wallets/groups directly. + return; + } + // Check if this account is already known by the tree to avoid double-insertion. if (!this.#accountIdToContext.has(account.id)) { this.update((state) => { - this.#insert(state.accountTree.wallets, account); + this.#insertAccount(state.accountTree.wallets, account); const context = this.#accountIdToContext.get(account.id); if (context) { @@ -851,50 +893,62 @@ export class AccountTreeController extends BaseController< if (context) { const { walletId, groupId } = context; - const previousSelectedAccountGroup = - this.state.accountTree.selectedAccountGroup; - let selectedAccountGroupChanged = false; + // If it's in the context, then `group` should be defined. + const group = this.state.accountTree.wallets[walletId]?.groups[groupId]; + assert( + group, + 'Account group associated with a context cannot be undefined', + ); - this.update((state) => { - const accounts = - state.accountTree.wallets[walletId]?.groups[groupId]?.accounts; + // We're only considering non-BIP-44 accounts for single account events. + if (group.type !== AccountGroupType.MultichainAccount) { + const index = group.accounts.indexOf(accountId); + + if (index !== -1) { + let selectedAccountGroupChanged = false; + const previousSelectedAccountGroup = + this.state.accountTree.selectedAccountGroup; + + this.update((state) => { + const accounts = + state.accountTree.wallets[walletId]?.groups[groupId].accounts; - if (accounts) { - const index = accounts.indexOf(accountId); - if (index !== -1) { + // Effectively remove account from the group and state. 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; + // If there's no more account, we have to prune the group and maybe select + // a new account group too. + if (accounts.length === 0) { + // Check if we need to update selectedAccountGroup after removal + if (state.accountTree.selectedAccountGroup === groupId) { + // 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; + } + + this.#pruneEmptyGroupAndWallet(state, walletId, groupId); } - } - if (accounts.length === 0) { - this.#pruneEmptyGroupAndWallet(state, walletId, groupId); + }); + + this.messenger.publish( + `${controllerName}:accountTreeChange`, + this.state.accountTree, + ); + + // Emit selectedAccountGroupChange event if the selected group changed + if (selectedAccountGroupChanged) { + this.messenger.publish( + `${controllerName}:selectedAccountGroupChange`, + this.state.accountTree.selectedAccountGroup, + previousSelectedAccountGroup, + ); } } - }); - this.messenger.publish( - `${controllerName}:accountTreeChange`, - this.state.accountTree, - ); - - // Emit selectedAccountGroupChange event if the selected group changed - if (selectedAccountGroupChanged) { - this.messenger.publish( - `${controllerName}:selectedAccountGroupChange`, - this.state.accountTree.selectedAccountGroup, - previousSelectedAccountGroup, - ); } // Clear reverse-mapping for that account. @@ -935,6 +989,26 @@ export class AccountTreeController extends BaseController< return state; } + /** + * Sorts accounts based on a pre-defined type orders. + * + * @param group Account group to sort. + */ + #sortAccountsOfGroup(group: AccountGroupObject) { + group.accounts.sort( + /* istanbul ignore next: Comparator branch execution (a===id vs b===id) + * and return attribution vary across engines; final ordering is covered + * by behavior tests. Ignoring the entire comparator avoids flaky line + * coverage without reducing scenario coverage. + */ + (a, b) => { + const aSortOrder = this.#accountIdToContext.get(a)?.sortOrder; + const bSortOrder = this.#accountIdToContext.get(b)?.sortOrder; + return (aSortOrder ?? MAX_SORT_ORDER) - (bSortOrder ?? MAX_SORT_ORDER); + }, + ); + } + /** * Insert an account inside an account tree. * @@ -945,12 +1019,18 @@ export class AccountTreeController extends BaseController< * @param wallets - Account tree. * @param account - The account to be inserted. */ - #insert( + #insertAccount( wallets: AccountTreeControllerState['accountTree']['wallets'], account: InternalAccount, ) { + // Those are owned by the `MultichainAccountService`, so they should be already handled + // by `#insertOrUpdateMultichainAccountWalletAndGroup` method. + assert( + !isBip44Account(account), + 'BIP-44 accounts cannot be inserted explicitly', + ); + const result = - this.#getEntropyRule().match(account) ?? this.#getSnapRule().match(account) ?? this.#getKeyringRule().match(account); // This one cannot fail. @@ -971,64 +1051,35 @@ export class AccountTreeController extends BaseController< // the union tag `result.wallet.type`. } as AccountWalletObject; wallet = wallets[walletId]; - - // Trigger atomic sync for new wallet (only for entropy wallets) - if (wallet.type === AccountWalletType.Entropy) { - this.#backupAndSyncService.enqueueSingleWalletSync(walletId); - } } const groupId = result.group.id; - let group = wallet.groups[groupId]; - const { type, id } = account; - const sortOrder = ACCOUNT_TYPE_TO_SORT_ORDER[type]; - - if (!group) { - log(`[${walletId}] Add new group: [${groupId}]`); - wallet.groups[groupId] = { - ...result.group, - // Type-wise, we are guaranteed to always have at least 1 account. - accounts: [id], - metadata: { - name: '', - ...{ pinned: false, hidden: false }, // Default UI states - ...result.group.metadata, // Allow rules to override defaults - }, - // We do need to type-cast since we're not narrowing `result` with - // the union tag `result.group.type`. - } as AccountGroupObject; - group = wallet.groups[groupId]; + assert( + result.group.type === AccountGroupType.SingleAccount, + `Account insertion should always result in a "${AccountGroupType.SingleAccount}" group type`, + ); + assert( + !wallet.groups[groupId], + 'Single account insertion cannot re-use existing group', + ); - // Map group ID to its containing wallet ID for efficient direct access - this.#groupIdToWalletId.set(groupId, walletId); + log(`[${walletId}] Add new group: [${groupId}]`); + const group = { + ...result.group, + // Type-wise, we are guaranteed to always have at least 1 account. + accounts: [account.id], + metadata: { + name: '', + ...{ pinned: false, hidden: false }, // Default UI states + ...result.group.metadata, // Allow rules to override defaults + }, + // We do need to type-cast since we're not narrowing `result` with + // the union tag `result.group.type`. + } as AccountGroupObject; + wallet.groups[groupId] = group; - // Trigger atomic sync for new group (only for entropy wallets) - if (wallet.type === AccountWalletType.Entropy) { - this.#backupAndSyncService.enqueueSingleGroupSync(groupId); - } - } else { - group.accounts.push(id); - // We need to do this at every insertion because race conditions can happen - // during the account creation process where one provider completes before the other. - // The discovery process in the service can also lead to some accounts being created "out of order". - const { accounts } = group; - accounts.sort( - /* istanbul ignore next: Comparator branch execution (a===id vs b===id) - * and return attribution vary across engines; final ordering is covered - * by behavior tests. Ignoring the entire comparator avoids flaky line - * coverage without reducing scenario coverage. - */ - (a, b) => { - const aSortOrder = - a === id ? sortOrder : this.#accountIdToContext.get(a)?.sortOrder; - const bSortOrder = - b === id ? sortOrder : this.#accountIdToContext.get(b)?.sortOrder; - return ( - (aSortOrder ?? MAX_SORT_ORDER) - (bSortOrder ?? MAX_SORT_ORDER) - ); - }, - ); - } + // Map group ID to its containing wallet ID for efficient direct access + this.#groupIdToWalletId.set(groupId, walletId); log( `[${groupId}] Add new account: { id: "${account.id}", type: "${account.type}", address: "${account.address}"`, ); @@ -1037,8 +1088,138 @@ export class AccountTreeController extends BaseController< this.#accountIdToContext.set(account.id, { walletId: wallet.id, groupId: group.id, - sortOrder, + sortOrder: ACCOUNT_TYPE_TO_SORT_ORDER[account.type], }); + + // We need to do this at every insertion because race conditions can happen + // during the account creation process where one provider completes before the other. + // The discovery process in the service can also lead to some accounts being created "out of order". + this.#sortAccountsOfGroup(group); + } + + /** + * Insert an account group inside an account tree. + * + * We go over multiple rules to try to "match" the account following + * specific criterias. If a rule "matches" an account, then this + * account get added into its proper account wallet and account group. + * + * @param wallets - Account tree. + * @param multichainAccountWallet - The multichain account wallet to be inserted or updated. + * @param multichainAccountGroup - The multichain account group to be inserted or updated. + */ + #insertOrUpdateMultichainAccountWalletAndGroup( + wallets: AccountTreeControllerState['accountTree']['wallets'], + multichainAccountWallet: MultichainAccountWallet< + Bip44Account + >, + multichainAccountGroup: MultichainAccountGroup< + Bip44Account + >, + ) { + const walletId = multichainAccountWallet.id; + const groupId = multichainAccountGroup.id; + + let wallet: AccountWalletEntropyObject | null = null; + let group: AccountGroupMultichainAccountObject | null = null; + + // Check type, mainly to infer proper wallet object type. + const walletObject = wallets[walletId]; + if (walletObject && walletObject.type === AccountWalletType.Entropy) { + wallet = walletObject; + group = wallet.groups[groupId]; + } + + // We always re-use the account list from the group (if accounts get + // added/removed/re-ordered within the group). + const accounts = multichainAccountGroup.getAccounts(); + const accountIds = accounts + // For now, we need this type-cast because `getAccounts` do not have the same + // type-constraint (uses string[] instead of [string; ...string]) + .map((account) => account.id) as AccountGroupObject['accounts']; + + // Create the group object first, to inject it in the wallet in case this wallet is + // not part of the tree yet. + if (!group) { + group = { + id: multichainAccountGroup.id, + type: multichainAccountGroup.type, + accounts: accountIds, + metadata: { + entropy: { + groupIndex: multichainAccountGroup.groupIndex, + }, + name: '', + pinned: false, + hidden: false, + }, + }; + } else { + // We clear existing contexts for previous accounts of that group because: + // - Accounts might have been removed + // - Accounts context mapping will get updated below + group.accounts.forEach((accountId) => + this.#accountIdToContext.delete(accountId), + ); + + group.accounts = accountIds; + } + + if (!wallet) { + wallet = { + id: multichainAccountWallet.id, + type: multichainAccountWallet.type, + status: multichainAccountWallet.status, + groups: { + [group.id]: group, + }, + metadata: { + entropy: { + id: multichainAccountWallet.entropySource, + }, + name: '', // Will get updated later. + }, + }; + wallets[walletId] = wallet; + + // Trigger atomic sync for new wallet. + this.#backupAndSyncService.enqueueSingleWalletSync(walletId); + } else { + wallet.groups[group.id] = group; + + // Trigger atomic sync for new group. + this.#backupAndSyncService.enqueueSingleGroupSync(groupId); + } + + // Map group ID to its containing wallet ID for efficient direct access + this.#groupIdToWalletId.set(groupId, walletId); + + // Update the reverse mapping for all accounts account. + for (const account of accounts) { + this.#accountIdToContext.set(account.id, { + walletId: wallet.id, + groupId: group.id, + sortOrder: ACCOUNT_TYPE_TO_SORT_ORDER[account.type], + }); + } + + // We need to do this at every insertion because race conditions can happen + // during the account creation process where one provider completes before the other. + // The discovery process in the service can also lead to some accounts being created "out of order". + this.#sortAccountsOfGroup(group); + } + + /** + * List all multichain accounts. + * + * @returns The list of all internal accounts. + */ + #getMultichainAccountWallets(): MultichainAccountWallet< + Bip44Account + >[] { + return this.messenger.call( + 'MultichainAccountService:getMultichainAccountWallets', + ); } /** @@ -1199,6 +1380,48 @@ export class AccountTreeController extends BaseController< ); } + /** + * Handles multichain account group created/updated event from + * the MultichainAccountService. + * + * @param multichainAccountGroup - Multichain account group being that got created or updated. + */ + #handleMultichainAccountGroupCreatedOrUpdated( + multichainAccountGroup: MultichainAccountGroup< + Bip44Account + >, + ): void { + // We wait for the first `init` to be called to actually build up the tree and + // mutate it. We expect the caller to first update the `AccountsController` state + // to force the migration of accounts, and then call `init`. + if (!this.#initialized) { + return; + } + + this.update((state) => { + this.#insertOrUpdateMultichainAccountWalletAndGroup( + state.accountTree.wallets, + multichainAccountGroup.wallet, + multichainAccountGroup, + ); + + const wallet = + state.accountTree.wallets[multichainAccountGroup.wallet.id]; + if (wallet) { + this.#applyAccountWalletMetadata(state, wallet.id); + + const group = wallet.groups[multichainAccountGroup.id]; + if (group) { + this.#applyAccountGroupMetadata(state, wallet.id, group.id); + } + } + }); + this.messenger.publish( + `${controllerName}:accountTreeChange`, + this.state.accountTree, + ); + } + /** * Handles multichain account wallet status change from * the MultichainAccountService. diff --git a/packages/account-tree-controller/src/rules/entropy.test.ts b/packages/account-tree-controller/src/rules/entropy.test.ts index 9023d344de4..dd20f33ac01 100644 --- a/packages/account-tree-controller/src/rules/entropy.test.ts +++ b/packages/account-tree-controller/src/rules/entropy.test.ts @@ -13,9 +13,11 @@ import { } from '@metamask/keyring-api'; import { KeyringTypes } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; -import type { AccountWalletEntropyObject } from 'src/wallet'; -import { EntropyRule } from './entropy'; +import { + getEntropyComputedAccountGroupName, + getEntropyDefaultAccountGroupPrefix, +} from './entropy'; import { getAccountTreeControllerMessenger, getRootMessenger, @@ -60,15 +62,13 @@ const MOCK_HD_ACCOUNT_1: Bip44Account = { }, }; -describe('EntropyRule', () => { - describe('getComputedAccountGroupName', () => { +describe('entropy', () => { + describe('getEntropyComputedAccountGroupName', () => { it('uses BaseRule implementation', () => { - const messenger = getRootMessenger(); - const accountTreeControllerMessenger = - getAccountTreeControllerMessenger(messenger); - const rule = new EntropyRule(accountTreeControllerMessenger); + const rootMessenger = getRootMessenger(); + const messenger = getAccountTreeControllerMessenger(rootMessenger); - messenger.registerActionHandler( + rootMessenger.registerActionHandler( 'AccountsController:getAccount', () => MOCK_HD_ACCOUNT_1, ); @@ -90,18 +90,16 @@ describe('EntropyRule', () => { }, }; - expect(rule.getComputedAccountGroupName(group)).toBe( + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe( MOCK_HD_ACCOUNT_1.metadata.name, ); }); it('returns empty string when account is not found', () => { - const messenger = getRootMessenger(); - const accountTreeControllerMessenger = - getAccountTreeControllerMessenger(messenger); - const rule = new EntropyRule(accountTreeControllerMessenger); + const rootMessenger = getRootMessenger(); + const messenger = getAccountTreeControllerMessenger(rootMessenger); - messenger.registerActionHandler( + rootMessenger.registerActionHandler( 'AccountsController:getAccount', () => undefined, ); @@ -123,26 +121,18 @@ describe('EntropyRule', () => { }, }; - expect(rule.getComputedAccountGroupName(group)).toBe(''); + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe(''); }); }); - describe('getDefaultAccountGroupPrefix', () => { + describe('getEntropyDefaultAccountGroupPrefix', () => { it('returns formatted account name prefix', () => { - const rootMessenger = getRootMessenger(); - const messenger = getAccountTreeControllerMessenger(rootMessenger); - const rule = new EntropyRule(messenger); - // The entropy wallet object is not used here. - const wallet = {} as unknown as AccountWalletEntropyObject; - - expect(rule.getDefaultAccountGroupPrefix(wallet)).toBe('Account'); + expect(getEntropyDefaultAccountGroupPrefix()).toBe('Account'); }); - it('getComputedAccountGroupName returns account name with EVM priority', () => { - const messenger = getRootMessenger(); - const accountTreeControllerMessenger = - getAccountTreeControllerMessenger(messenger); - const rule = new EntropyRule(accountTreeControllerMessenger); + it('getEntropyComputedAccountGroupName returns account name with EVM priority', () => { + const rootMessenger = getRootMessenger(); + const messenger = getAccountTreeControllerMessenger(rootMessenger); const mockEvmAccount: InternalAccount = { ...MOCK_HD_ACCOUNT_1, @@ -154,7 +144,7 @@ describe('EntropyRule', () => { }, }; - messenger.registerActionHandler( + rootMessenger.registerActionHandler( 'AccountsController:getAccount', () => mockEvmAccount, ); @@ -176,16 +166,16 @@ describe('EntropyRule', () => { }, }; - expect(rule.getComputedAccountGroupName(group)).toBe('EVM Account'); + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe( + 'EVM Account', + ); }); - it('getComputedAccountGroupName returns empty string when no accounts found', () => { - const messenger = getRootMessenger(); - const accountTreeControllerMessenger = - getAccountTreeControllerMessenger(messenger); - const rule = new EntropyRule(accountTreeControllerMessenger); + it('getEntropyComputedAccountGroupName returns empty string when no accounts found', () => { + const rootMessenger = getRootMessenger(); + const messenger = getAccountTreeControllerMessenger(rootMessenger); - messenger.registerActionHandler( + rootMessenger.registerActionHandler( 'AccountsController:getAccount', () => undefined, ); @@ -207,14 +197,12 @@ describe('EntropyRule', () => { }, }; - expect(rule.getComputedAccountGroupName(group)).toBe(''); + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe(''); }); - it('getComputedAccountGroupName returns empty string for non-EVM accounts to prevent chain-specific names', () => { - const messenger = getRootMessenger(); - const accountTreeControllerMessenger = - getAccountTreeControllerMessenger(messenger); - const rule = new EntropyRule(accountTreeControllerMessenger); + it('getEntropyComputedAccountGroupName returns empty string for non-EVM accounts to prevent chain-specific names', () => { + const rootMessenger = getRootMessenger(); + const messenger = getAccountTreeControllerMessenger(rootMessenger); const mockSolanaAccount: InternalAccount = { ...MOCK_HD_ACCOUNT_1, @@ -226,7 +214,7 @@ describe('EntropyRule', () => { }, }; - messenger.registerActionHandler( + rootMessenger.registerActionHandler( 'AccountsController:getAccount', (accountId: string) => { const accounts: Record = { @@ -254,13 +242,12 @@ describe('EntropyRule', () => { }; // Should return empty string, not "Solana Account 2", to fallback to default naming - expect(rule.getComputedAccountGroupName(group)).toBe(''); + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe(''); }); it('getComputedAccountGroupName returns EVM name even when non-EVM accounts are present first', () => { const rootMessenger = getRootMessenger(); const messenger = getAccountTreeControllerMessenger(rootMessenger); - const rule = new EntropyRule(messenger); const mockSolanaAccount: InternalAccount = { ...MOCK_HD_ACCOUNT_1, @@ -311,7 +298,9 @@ describe('EntropyRule', () => { }; // Should return EVM account name, not Solana account name - expect(rule.getComputedAccountGroupName(group)).toBe('Main Account'); + expect(getEntropyComputedAccountGroupName(messenger, group)).toBe( + 'Main Account', + ); }); }); }); diff --git a/packages/account-tree-controller/src/rules/entropy.ts b/packages/account-tree-controller/src/rules/entropy.ts index 44b944bae68..44ab2a043a8 100644 --- a/packages/account-tree-controller/src/rules/entropy.ts +++ b/packages/account-tree-controller/src/rules/entropy.ts @@ -1,114 +1,82 @@ -import { +import type { AccountGroupType, AccountWalletType, - isBip44Account, - toMultichainAccountGroupId, - toMultichainAccountWalletId, } from '@metamask/account-api'; import { isEvmAccountType } from '@metamask/keyring-api'; import { KeyringTypes } from '@metamask/keyring-controller'; -import type { InternalAccount } from '@metamask/keyring-internal-api'; import type { AccountGroupObjectOf } from '../group'; -import { BaseRule, type Rule, type RuleResult } from '../rule'; +import type { AccountTreeControllerMessenger } from '../types'; import type { AccountWalletObjectOf } from '../wallet'; -export class EntropyRule - extends BaseRule - implements Rule -{ - readonly walletType = AccountWalletType.Entropy; - - readonly groupType = AccountGroupType.MultichainAccount; - - getEntropySourceIndex(entropySource: string) { - const { keyrings } = this.messenger.call('KeyringController:getState'); - - return keyrings - .filter((keyring) => keyring.type === (KeyringTypes.hd as string)) - .findIndex((keyring) => keyring.metadata.id === entropySource); - } +/** + * Get the entropy index for a given entropy source. + * + * @param messenger - Controller's messenger. + * @param entropySource - The entropy source to get the index for. + * @returns The index for this entropy source. + */ +function getEntropySourceIndex( + messenger: AccountTreeControllerMessenger, + entropySource: string, +) { + const { keyrings } = messenger.call('KeyringController:getState'); + + return keyrings + .filter((keyring) => keyring.type === (KeyringTypes.hd as string)) + .findIndex((keyring) => keyring.metadata.id === entropySource); +} - match( - account: InternalAccount, - ): - | RuleResult - | undefined { - if (!isBip44Account(account)) { - return undefined; - } +/** + * Get default wallet name for entropy-based wallets. + * + * @param messenger - Controller's messenger. + * @param wallet - The wallet to get the name for. + * @returns The default name for this wallet. + */ +export function getEntropyDefaultAccountWalletName( + messenger: AccountTreeControllerMessenger, + wallet: AccountWalletObjectOf, +): string { + // NOTE: We have checked during the rule matching, so we can safely assume it will + // well-defined here. + const entropySourceIndex = getEntropySourceIndex( + messenger, + wallet.metadata.entropy.id, + ); + + return `Wallet ${entropySourceIndex + 1}`; // Use human indexing (starts at 1). +} - const entropySource = account.options.entropy.id; - const entropySourceIndex = this.getEntropySourceIndex(entropySource); - if (entropySourceIndex === -1) { - console.warn( - `! Found an unknown entropy ID: "${entropySource}", account "${account.id}" won't be grouped by entropy.`, - ); - return undefined; +/** + * Get computed group name for entropy-based wallet groups. + * + * @param messenger - Controller's messenger. + * @param group - The group to get the name for. + * @returns The computed name for this group. + */ +export function getEntropyComputedAccountGroupName( + messenger: AccountTreeControllerMessenger, + group: AccountGroupObjectOf, +): string { + // Only use EVM account names for multichain groups to avoid chain-specific names becoming group names. + // Non-EVM account names should not be used as group names since groups represent multichain collections. + for (const id of group.accounts) { + const account = messenger.call('AccountsController:getAccount', id); + + if (account && isEvmAccountType(account.type)) { + return account.metadata.name; } - - const walletId = toMultichainAccountWalletId(entropySource); - const groupId = toMultichainAccountGroupId( - walletId, - account.options.entropy.groupIndex, - ); - - return { - wallet: { - type: this.walletType, - id: walletId, - metadata: { - entropy: { - id: entropySource, - }, - }, - }, - - group: { - type: this.groupType, - id: groupId, - metadata: { - entropy: { - groupIndex: account.options.entropy.groupIndex, - }, - pinned: false, - hidden: false, - }, - }, - }; } - getDefaultAccountWalletName( - wallet: AccountWalletObjectOf, - ): string { - // NOTE: We have checked during the rule matching, so we can safely assume it will - // well-defined here. - const entropySourceIndex = this.getEntropySourceIndex( - wallet.metadata.entropy.id, - ); - - return `Wallet ${entropySourceIndex + 1}`; // Use human indexing (starts at 1). - } - - getComputedAccountGroupName( - group: AccountGroupObjectOf, - ): string { - // Only use EVM account names for multichain groups to avoid chain-specific names becoming group names. - // Non-EVM account names should not be used as group names since groups represent multichain collections. - for (const id of group.accounts) { - const account = this.messenger.call('AccountsController:getAccount', id); - - if (account && isEvmAccountType(account.type)) { - return account.metadata.name; - } - } - - return ''; - } + return ''; +} - getDefaultAccountGroupPrefix( - _wallet: AccountWalletObjectOf, - ): string { - return 'Account'; - } +/** + * Get the group name prefix for entropy-based wallet groups. + * + * @returns The prefix for entropy-based wallet groups. + */ +export function getEntropyDefaultAccountGroupPrefix(): string { + return 'Account'; } diff --git a/packages/account-tree-controller/src/types.ts b/packages/account-tree-controller/src/types.ts index da288a1b565..f443d43bb23 100644 --- a/packages/account-tree-controller/src/types.ts +++ b/packages/account-tree-controller/src/types.ts @@ -39,7 +39,12 @@ import type { AccountWalletObject, AccountTreeWalletPersistedMetadata, } from './wallet'; -import type { MultichainAccountServiceWalletStatusChangeEvent } from '../../multichain-account-service/src/types'; +import type { + MultichainAccountServiceGetMultichainAccountWalletsAction, + MultichainAccountServiceMultichainAccountGroupCreatedEvent, + MultichainAccountServiceMultichainAccountGroupUpdatedEvent, + MultichainAccountServiceWalletStatusChangeEvent, +} from '../../multichain-account-service/src/types'; // Backward compatibility aliases using indexed access types /** @@ -127,6 +132,7 @@ export type AllowedActions = | UserStorageController.UserStorageControllerPerformSetStorage | UserStorageController.UserStorageControllerPerformBatchSetStorage | AuthenticationController.AuthenticationControllerGetSessionProfile + | MultichainAccountServiceGetMultichainAccountWalletsAction | MultichainAccountServiceCreateMultichainAccountGroupAction; export type AccountTreeControllerActions = @@ -167,6 +173,8 @@ export type AllowedEvents = | AccountsControllerAccountRemovedEvent | AccountsControllerSelectedAccountChangeEvent | UserStorageController.UserStorageControllerStateChangeEvent + | MultichainAccountServiceMultichainAccountGroupCreatedEvent + | MultichainAccountServiceMultichainAccountGroupUpdatedEvent | MultichainAccountServiceWalletStatusChangeEvent; export type AccountTreeControllerEvents = diff --git a/packages/account-tree-controller/tests/mockMessenger.ts b/packages/account-tree-controller/tests/mockMessenger.ts index 10fa770a675..4dfa8081d3f 100644 --- a/packages/account-tree-controller/tests/mockMessenger.ts +++ b/packages/account-tree-controller/tests/mockMessenger.ts @@ -50,6 +50,8 @@ export function getAccountTreeControllerMessenger( 'AccountsController:selectedAccountChange', 'UserStorageController:stateChange', 'MultichainAccountService:walletStatusChange', + 'MultichainAccountService:multichainAccountGroupCreated', + 'MultichainAccountService:multichainAccountGroupUpdated', ], actions: [ 'AccountsController:listMultichainAccounts', @@ -63,6 +65,7 @@ export function getAccountTreeControllerMessenger( 'UserStorageController:performBatchSetStorage', 'AuthenticationController:getSessionProfile', 'MultichainAccountService:createMultichainAccountGroup', + 'MultichainAccountService:getMultichainAccountWallets', 'KeyringController:getState', 'SnapController:get', ],