diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index a6108aa0b1..27ec29543d 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Fetch account balances via RPC when switching EVM accounts, enabling RPC-only networks, switching account groups, or after a new account is added to the account tree ([#9388](https://github.com/MetaMask/core/pull/9388)) + - Integrators must delegate `AccountsController:selectedEvmAccountChange` to the AssetsController messenger + ## [10.0.1] ### Changed diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index cd16591df8..4a0da9c85d 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -2,8 +2,12 @@ import type { AccountTreeControllerGetAccountsFromSelectedAccountGroupAction, AccountTreeControllerSelectedAccountGroupChangeEvent, AccountTreeControllerState, + AccountTreeControllerStateChangeEvent, } from '@metamask/account-tree-controller'; -import type { AccountsControllerGetSelectedAccountAction } from '@metamask/accounts-controller'; +import type { + AccountsControllerGetSelectedAccountAction, + AccountsControllerSelectedEvmAccountChangeEvent, +} from '@metamask/accounts-controller'; import { BaseController } from '@metamask/base-controller'; import type { ControllerGetStateAction, @@ -11,7 +15,10 @@ import type { ControllerStateChangedEvent, StateMetadata, } from '@metamask/base-controller'; -import type { ClientControllerState } from '@metamask/client-controller'; +import type { + ClientControllerState, + ClientControllerStateChangeEvent, +} from '@metamask/client-controller'; import { clientControllerSelectors } from '@metamask/client-controller'; import type { TraceCallback } from '@metamask/controller-utils'; import type { @@ -39,8 +46,8 @@ import type { } from '@metamask/network-controller'; import type { NetworkEnablementControllerGetStateAction, - NetworkEnablementControllerEvents, NetworkEnablementControllerState, + NetworkEnablementControllerStateChangeEvent, } from '@metamask/network-enablement-controller'; import type { GetPermissions, @@ -321,26 +328,16 @@ type AllowedActions = // PhishingController | PhishingControllerBulkScanTokensAction; -type AccountTreeControllerStateChangedEvent = ControllerStateChangedEvent< - 'AccountTreeController', - AccountTreeControllerState ->; - -type ClientControllerStateChangedEvent = ControllerStateChangedEvent< - 'ClientController', - ClientControllerState ->; - -type NetworkEnablementControllerStateChangedEvent = ControllerStateChangedEvent< - 'NetworkEnablementController', - NetworkEnablementControllerState ->; - type AllowedEvents = // AssetsController | AccountTreeControllerSelectedAccountGroupChangeEvent - | AccountTreeControllerStateChangedEvent - | ClientControllerStateChangedEvent + | AccountTreeControllerStateChangeEvent + | ControllerStateChangedEvent< + 'AccountTreeController', + AccountTreeControllerState + > + | ClientControllerStateChangeEvent + | ControllerStateChangedEvent<'ClientController', ClientControllerState> | KeyringControllerLockEvent | KeyringControllerUnlockEvent | PreferencesControllerStateChangeEvent @@ -355,10 +352,14 @@ type AllowedEvents = | NetworkControllerNetworkDidChangeEvent | NetworkControllerNetworkRemovedEvent // StakedBalanceDataSource - | NetworkEnablementControllerEvents - | NetworkEnablementControllerStateChangedEvent + | NetworkEnablementControllerStateChangeEvent + | ControllerStateChangedEvent< + 'NetworkEnablementController', + NetworkEnablementControllerState + > // SnapDataSource | AccountsControllerAccountBalancesUpdatedEvent + | AccountsControllerSelectedEvmAccountChangeEvent | PermissionControllerStateChange | SnapControllerSnapInstalledEvent // BackendWebsocketDataSource @@ -1079,6 +1080,14 @@ export class AssetsController extends BaseController< }, ); + // Switching EVM address within the same account group (e.g. Account 1 → Account 2). + this.messenger.subscribe( + 'AccountsController:selectedEvmAccountChange', + (account) => { + this.#handleSelectedEvmAccountChanged(account).catch(console.error); + }, + ); + // Catch the initial tree build. On returning users, // `selectedAccountGroupChange` does NOT fire when the persisted group // is unchanged, and `accountTreeChange` doesn't fire either (init() @@ -1329,10 +1338,14 @@ export class AssetsController extends BaseController< currentCount: currentIds.size, }); + const newAccounts = accounts.filter( + (account) => !this.#lastKnownAccountIds.has(account.id), + ); + this.#lastKnownAccountIds = currentIds; this.#ensureNativeBalancesDefaultZero(); this.#ensureDefaultTrackedAssetsSeeded(); - this.#runAccountTreeRefresh(accounts).catch((error) => { + this.#runAccountTreeRefresh(accounts, newAccounts).catch((error) => { log('Failed to refresh assets after tree change', error); }); } else { @@ -1340,7 +1353,10 @@ export class AssetsController extends BaseController< } } - async #runAccountTreeRefresh(accounts: InternalAccount[]): Promise { + async #runAccountTreeRefresh( + accounts: InternalAccount[], + newAccounts: InternalAccount[] = [], + ): Promise { const releaseLock = await this.#accountRefreshMutex.acquire(); try { await this.getAssets(accounts, { @@ -1348,6 +1364,11 @@ export class AssetsController extends BaseController< forceUpdate: true, }); this.#subscribeAssets(); + if (newAccounts.length > 0) { + await this.#fetchAccountBalancesViaRpc(newAccounts, [ + ...this.#enabledChains, + ]); + } } catch (error) { log('Failed to fetch assets after tree change', error); this.#subscribeAssets(); @@ -1356,6 +1377,42 @@ export class AssetsController extends BaseController< } } + /** + * Fetch balances via RpcDataSource and merge into state. + * + * @param accounts - Accounts to fetch for. + * @param chainIds - Chains to fetch on. + */ + async #fetchAccountBalancesViaRpc( + accounts: InternalAccount[], + chainIds: ChainId[], + ): Promise { + if (accounts.length === 0 || chainIds.length === 0) { + return; + } + + const rpcChains = chainIds.filter((chainId) => + this.#rpcDataSource.getActiveChainsSync().includes(chainId), + ); + if (rpcChains.length === 0) { + return; + } + + const request = this.#buildDataRequest(accounts, rpcChains, { + dataTypes: ['balance'], + forceUpdate: true, + }); + const { response } = await this.#executeMiddlewares( + [ + createParallelBalanceMiddleware([this.#rpcDataSource]), + this.#detectionMiddleware, + ], + request, + ); + + await this.#updateState({ ...response, updateMode: 'merge' }); + } + /** * Force-fetch balances then subscribe. Used on unlock / first startup. * @@ -3320,6 +3377,10 @@ export class AssetsController extends BaseController< // Subscribe after fetch so WS notifications can recover state this.#subscribeAssets(); + await this.#fetchAccountBalancesViaRpc(accounts, [ + ...this.#enabledChains, + ]); + this.#ensureNativeBalancesDefaultZero(); this.#ensureDefaultTrackedAssetsSeeded(); } finally { @@ -3327,6 +3388,40 @@ export class AssetsController extends BaseController< } } + /** + * Handle EVM account selection within the current account group. + * Fires when the user picks a different address under the same logical + * account (not when switching account groups). + * + * @param account - Newly selected EVM account. + */ + async #handleSelectedEvmAccountChanged( + account: InternalAccount, + ): Promise { + if (!this.#uiOpen || !this.#keyringUnlocked || !this.#isEnabled()) { + return; + } + + const releaseLock = await this.#accountRefreshMutex.acquire(); + try { + await this.getAssets([account], { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }); + + this.#subscribeAssets(); + + await this.#fetchAccountBalancesViaRpc( + [account], + [...this.#enabledChains], + ); + + this.#ensureNativeBalancesDefaultZero(); + } finally { + releaseLock(); + } + } + async #handleEnabledNetworksChanged( enabledNetworkMap: NetworkEnablementControllerState['enabledNetworkMap'], ): Promise { @@ -3364,12 +3459,15 @@ export class AssetsController extends BaseController< this.#subscribeAssets(); // Do one-time fetch for newly enabled chains; merge so we keep existing chain balances - if (addedChains.length > 0 && this.#getSelectedAccounts().length > 0) { - await this.getAssets(this.#getSelectedAccounts(), { + const accounts = this.#getSelectedAccounts(); + if (addedChains.length > 0 && accounts.length > 0) { + await this.getAssets(accounts, { chainIds: addedChains, forceUpdate: true, updateMode: 'merge', }); + + await this.#fetchAccountBalancesViaRpc(accounts, addedChains); } this.#ensureNativeBalancesDefaultZero(); @@ -3487,6 +3585,8 @@ export class AssetsController extends BaseController< dataTypes: ['balance', 'metadata', 'price'], }); + await this.#fetchAccountBalancesViaRpc(accounts, [selectedChainId]); + this.#ensureNativeBalancesDefaultZero(); this.#fetchMissingPricesWithoutCache(accounts, [selectedChainId]); } finally {