fix(assets-controller): seed native balances after account refresh

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

packages/assets-controller/src/AssetsController.ts

@@ -2,16 +2,23 @@ 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,
   ControllerStateChangeEvent,
   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,25 +1338,37 @@ 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 {
       this.#start();
     }
   }
 
-  async #runAccountTreeRefresh(accounts: InternalAccount[]): Promise<void> {
+  async #runAccountTreeRefresh(
+    accounts: InternalAccount[],
+    newAccounts: InternalAccount[] = [],
+  ): Promise<void> {
     const releaseLock = await this.#accountRefreshMutex.acquire();
     try {
       await this.getAssets(accounts, {
         chainIds: [...this.#enabledChains],
         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<void> {
+    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,13 +3377,51 @@ 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 {
       releaseLock();
     }
   }
 
+  /**
+   * 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<void> {
+    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<void> {
@@ -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 {