From ad2c7e232728d360c1839b78609a540c555ad24f Mon Sep 17 00:00:00 2001 From: gantunesr <17601467+gantunesr@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:50:59 -0400 Subject: [PATCH] feat: add getAccountFromSelectedAccountGroup method --- ...countTreeController-method-action-types.ts | 25 ++ .../src/AccountTreeController.test.ts | 246 ++++++++++++++++++ .../src/AccountTreeController.ts | 76 +++++- packages/account-tree-controller/src/index.ts | 1 + 4 files changed, 335 insertions(+), 13 deletions(-) diff --git a/packages/account-tree-controller/src/AccountTreeController-method-action-types.ts b/packages/account-tree-controller/src/AccountTreeController-method-action-types.ts index abc6ca61e9d..8750542433c 100644 --- a/packages/account-tree-controller/src/AccountTreeController-method-action-types.ts +++ b/packages/account-tree-controller/src/AccountTreeController-method-action-types.ts @@ -65,6 +65,30 @@ export type AccountTreeControllerGetAccountsFromSelectedAccountGroupAction = { handler: AccountTreeController['getAccountsFromSelectedAccountGroup']; }; +/** + * Gets an account from the currently selected account group, optionally + * filtered by a CAIP-2 chain ID. + * + * This is the group-based replacement for both + * `AccountsController:getSelectedAccount` and + * `AccountsController:getSelectedMultichainAccount`. + * + * When no chain ID is provided, an account of the selected group is returned + * using an EVM-priority rule: the first EVM account found in the group, or the + * first account in the group if no EVM account is found. When a chain ID is + * provided, the first account in the selected group whose scopes match the + * given chain is returned. + * + * @param chainId - Optional CAIP-2 chain ID used to filter accounts by scope. + * @returns The matching internal account from the selected group, or + * undefined if no group is selected or no account matches. + * @throws If `chainId` is provided but is not a valid CAIP-2 chain ID. + */ +export type AccountTreeControllerGetAccountFromSelectedAccountGroupAction = { + type: `AccountTreeController:getAccountFromSelectedAccountGroup`; + handler: AccountTreeController['getAccountFromSelectedAccountGroup']; +}; + /** * Gets the account group object from its ID. * @@ -208,6 +232,7 @@ export type AccountTreeControllerMethodActions = | AccountTreeControllerGetAccountWalletObjectAction | AccountTreeControllerGetAccountWalletObjectsAction | AccountTreeControllerGetAccountsFromSelectedAccountGroupAction + | AccountTreeControllerGetAccountFromSelectedAccountGroupAction | AccountTreeControllerGetAccountGroupObjectAction | AccountTreeControllerGetAccountContextAction | AccountTreeControllerGetSelectedAccountGroupAction diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index dfdf7bfa5e8..1916481eb95 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -30,6 +30,7 @@ import type { KeyringObject } from '@metamask/keyring-controller'; import { KeyringTypes } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import type { GetSnap as SnapControllerGetSnap } from '@metamask/snaps-controllers'; +import type { CaipChainId } from '@metamask/utils'; import { getAccountTreeControllerMessenger, @@ -1230,6 +1231,212 @@ describe('AccountTreeController', () => { }); }); + describe('getAccountFromSelectedAccountGroup', () => { + const evmAndNonEvmGroupSetup = (): { + controller: AccountTreeController; + mockSolAccount1: Bip44Account; + } => { + const mockSolAccount1: Bip44Account = { + ...MOCK_SNAP_ACCOUNT_1, + options: { + entropy: { + ...MOCK_SNAP_ACCOUNT_1.options.entropy, + groupIndex: 0, + }, + }, + }; + + const { controller } = setup({ + accounts: [MOCK_HD_ACCOUNT_2, mockSolAccount1], + keyrings: [MOCK_HD_KEYRING_2], + }); + + controller.init(); + + return { controller, mockSolAccount1 }; + }; + + it('returns the EVM account when the selected group has both EVM and non-EVM accounts', () => { + const { controller } = evmAndNonEvmGroupSetup(); + + expect( + controller.getAccountFromSelectedAccountGroup(), + ).toStrictEqual(MOCK_HD_ACCOUNT_2); + }); + + it('returns the first account when the selected group has no EVM account', () => { + const nonEvmGroupId = toAccountGroupId( + MOCK_PREPOPULATED_WALLET_ID, + 'tron-group', + ); + const nonEvmGroupState: Partial = { + selectedAccountGroup: nonEvmGroupId, + accountTree: { + wallets: { + [MOCK_PREPOPULATED_WALLET_ID]: { + id: MOCK_PREPOPULATED_WALLET_ID, + type: AccountWalletType.Entropy, + status: 'ready', + groups: { + [nonEvmGroupId]: { + id: nonEvmGroupId, + type: AccountGroupType.MultichainAccount, + accounts: [MOCK_TRX_ACCOUNT_1.id], + metadata: { + name: 'Tron Group', + entropy: { groupIndex: 0 }, + pinned: false, + hidden: false, + lastSelected: 0, + }, + }, + }, + metadata: { + name: 'Wallet 1', + entropy: { id: MOCK_HD_KEYRING_1.metadata.id }, + }, + }, + }, + }, + }; + + const { controller } = setup({ + accounts: [MOCK_TRX_ACCOUNT_1], + keyrings: [], + state: nonEvmGroupState, + }); + + expect( + controller.getAccountFromSelectedAccountGroup(), + ).toStrictEqual(MOCK_TRX_ACCOUNT_1); + }); + + it('returns undefined when no group is selected', () => { + const { controller } = setup({ + accounts: [], + keyrings: [], + }); + + controller.init(); + + expect( + controller.getAccountFromSelectedAccountGroup(), + ).toBeUndefined(); + }); + + it('returns undefined when the selected group is empty', () => { + const emptyGroupId = toAccountGroupId( + MOCK_PREPOPULATED_WALLET_ID, + 'empty-group', + ); + const emptyGroupState: Partial = { + selectedAccountGroup: emptyGroupId, + accountTree: { + wallets: { + [MOCK_PREPOPULATED_WALLET_ID]: { + id: MOCK_PREPOPULATED_WALLET_ID, + type: AccountWalletType.Entropy, + status: 'ready', + groups: { + [emptyGroupId]: { + id: emptyGroupId, + type: AccountGroupType.MultichainAccount, + accounts: [], + metadata: { + name: 'Empty Group', + entropy: { groupIndex: 0 }, + pinned: false, + hidden: false, + lastSelected: 0, + }, + }, + }, + metadata: { + name: 'Wallet 1', + entropy: { id: MOCK_HD_KEYRING_1.metadata.id }, + }, + }, + }, + }, + }; + + const { controller } = setup({ + accounts: [], + keyrings: [], + state: emptyGroupState, + }); + + expect( + controller.getAccountFromSelectedAccountGroup(), + ).toBeUndefined(); + }); + + it('returns the account 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 — the account must be resolved from persisted state. + expect( + controller.getAccountFromSelectedAccountGroup(), + ).toStrictEqual(MOCK_HD_ACCOUNT_1); + }); + + it('returns the account matching the given CAIP-2 chain ID (non-EVM)', () => { + const { controller, mockSolAccount1 } = evmAndNonEvmGroupSetup(); + + expect( + controller.getAccountFromSelectedAccountGroup(SolScope.Mainnet), + ).toStrictEqual(mockSolAccount1); + }); + + it('returns the account matching the given CAIP-2 chain ID (EVM)', () => { + const { controller } = evmAndNonEvmGroupSetup(); + + expect( + controller.getAccountFromSelectedAccountGroup(EthScope.Mainnet), + ).toStrictEqual(MOCK_HD_ACCOUNT_2); + }); + + it('returns undefined when no group is selected and a chain ID is provided', () => { + const { controller } = setup({ + accounts: [], + keyrings: [], + }); + + controller.init(); + + expect( + controller.getAccountFromSelectedAccountGroup(SolScope.Mainnet), + ).toBeUndefined(); + }); + + it('returns undefined when no account in the selected group matches the scope', () => { + const { controller } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + state: MOCK_PREPOPULATED_STATE, + }); + + // The selected group only contains an EVM account, so a Sol scope matches nothing. + expect( + controller.getAccountFromSelectedAccountGroup(SolScope.Mainnet), + ).toBeUndefined(); + }); + + it('throws when the chain ID is not a valid CAIP-2 chain ID', () => { + const { controller } = evmAndNonEvmGroupSetup(); + + expect(() => + controller.getAccountFromSelectedAccountGroup( + 'not-a-caip-chain-id' as CaipChainId, + ), + ).toThrow('Invalid CAIP-2 chain ID: not-a-caip-chain-id'); + }); + }); + describe('on AccountsController:accountsRemoved', () => { it('removes an account from the tree', () => { // 2 accounts that share the same entropy source (thus, same wallet). @@ -4285,6 +4492,45 @@ describe('AccountTreeController', () => { expect(spy).toHaveBeenCalled(); }); + it('calls getAccountFromSelectedAccountGroup via AccountTreeController:getAccountFromSelectedAccountGroup', () => { + const spy = jest.spyOn( + AccountTreeController.prototype, + 'getAccountFromSelectedAccountGroup', + ); + + const { controller, messenger } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + controller.init(); + + messenger.call( + 'AccountTreeController:getAccountFromSelectedAccountGroup', + ); + expect(spy).toHaveBeenCalled(); + }); + + it('calls getAccountFromSelectedAccountGroup with a chain ID via the messenger action', () => { + const spy = jest.spyOn( + AccountTreeController.prototype, + 'getAccountFromSelectedAccountGroup', + ); + + const { controller, messenger } = setup({ + accounts: [MOCK_HD_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + controller.init(); + + messenger.call( + 'AccountTreeController:getAccountFromSelectedAccountGroup', + EthScope.Mainnet as CaipChainId, + ); + expect(spy).toHaveBeenCalledWith(EthScope.Mainnet as CaipChainId); + }); + it('gets account context with AccountTreeController:getAccountContext', () => { const spy = jest.spyOn( AccountTreeController.prototype, diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index e37b2dc1524..7ef3b44973a 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -13,7 +13,8 @@ import { BaseController } from '@metamask/base-controller'; import type { TraceCallback } from '@metamask/controller-utils'; import { isEvmAccountType } from '@metamask/keyring-api'; import type { InternalAccount } from '@metamask/keyring-internal-api'; -import { assert } from '@metamask/utils'; +import { assert, isCaipChainId } from '@metamask/utils'; +import type { CaipChainId } from '@metamask/utils'; import type { BackupAndSyncEmitAnalyticsEventParams } from './backup-and-sync/analytics'; import { @@ -49,6 +50,7 @@ const MESSENGER_EXPOSED_METHODS = [ 'getSelectedAccountGroup', 'setSelectedAccountGroup', 'getAccountsFromSelectedAccountGroup', + 'getAccountFromSelectedAccountGroup', 'getAccountContext', 'setAccountWalletName', 'setAccountGroupName', @@ -858,6 +860,48 @@ export class AccountTreeController extends BaseController< return selector ? select(accounts, selector) : accounts; } + /** + * Gets an account from the currently selected account group, optionally + * filtered by a CAIP-2 chain ID. + * + * This is the group-based replacement for both + * `AccountsController:getSelectedAccount` and + * `AccountsController:getSelectedMultichainAccount`. + * + * When no chain ID is provided, an account of the selected group is returned + * using an EVM-priority rule: the first EVM account found in the group, or the + * first account in the group if no EVM account is found. When a chain ID is + * provided, the first account in the selected group whose scopes match the + * given chain is returned. + * + * @param chainId - Optional CAIP-2 chain ID used to filter accounts by scope. + * @returns The matching internal account from the selected group, or + * undefined if no group is selected or no account matches. + * @throws If `chainId` is provided but is not a valid CAIP-2 chain ID. + */ + getAccountFromSelectedAccountGroup( + chainId?: CaipChainId, + ): InternalAccount | undefined { + const groupId = this.getSelectedAccountGroup(); + if (!groupId) { + return undefined; + } + + if (!chainId) { + return this.#getAccountFromAccountGroupId(groupId); + } + + if (!isCaipChainId(chainId)) { + throw new Error(`Invalid CAIP-2 chain ID: ${String(chainId)}`); + } + + const accounts = this.getAccountsFromSelectedAccountGroup({ + scopes: [chainId], + }); + + return accounts[0]; + } + /** * Gets the account group object from its ID. * @@ -1303,7 +1347,7 @@ export class AccountTreeController extends BaseController< } // Find the first account in this group to select - const accountToSelect = this.#getDefaultAccountFromAccountGroupId(groupId); + const accountToSelect = this.#getAccountFromAccountGroupId(groupId); if (!accountToSelect) { throw new Error('No accounts found in group'); } @@ -1343,10 +1387,10 @@ export class AccountTreeController extends BaseController< // but our handler is idempotent so it won't cause infinite loop this.messenger.call( 'AccountsController:setSelectedAccount', - accountToSelect, + accountToSelect.id, ); - log(`Selected account is now: ${accountToSelect}`); + log(`Selected account is now: ${accountToSelect.id}`); } } @@ -1439,35 +1483,41 @@ export class AccountTreeController extends BaseController< } /** - * Gets the default account for specified group. + * Gets an account for the specified group using the EVM-priority rule. + * + * The account is the first EVM account found in the group, or the first + * account in the group if no EVM account is found. * * @param groupId - The account group ID. - * @returns The first account ID in the group, or undefined if no accounts found. + * @returns The internal account in the group, or undefined if the group does + * not exist or is empty. */ - #getDefaultAccountFromAccountGroupId( + #getAccountFromAccountGroupId( groupId: AccountGroupId, - ): AccountId | undefined { + ): InternalAccount | undefined { const group = this.#getAccountGroup(groupId); if (group) { - let candidate; + let candidateId: AccountId | undefined; for (const id of group.accounts) { const account = this.messenger.call( 'AccountsController:getAccount', id, ); - if (!candidate) { - candidate = id; + if (!candidateId) { + candidateId = id; } if (account && isEvmAccountType(account.type)) { // EVM accounts have a higher priority, so if we find any, we just // use that account! - return account.id; + return account; } } - return candidate; + return candidateId + ? this.messenger.call('AccountsController:getAccount', candidateId) + : undefined; } return undefined; diff --git a/packages/account-tree-controller/src/index.ts b/packages/account-tree-controller/src/index.ts index 28b2f17f645..454cd53568e 100644 --- a/packages/account-tree-controller/src/index.ts +++ b/packages/account-tree-controller/src/index.ts @@ -25,6 +25,7 @@ export type { AccountTreeControllerGetAccountWalletObjectAction, AccountTreeControllerGetAccountWalletObjectsAction, AccountTreeControllerGetAccountsFromSelectedAccountGroupAction, + AccountTreeControllerGetAccountFromSelectedAccountGroupAction, AccountTreeControllerGetAccountGroupObjectAction, AccountTreeControllerGetAccountContextAction, AccountTreeControllerGetSelectedAccountGroupAction,