feat: changes to AccountTrackerController to enable migration in extension

packages/assets-controllers/CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- **BREAKING:** Modify AccountTrackerController events and actions ([#6938](https://github.com/MetaMask/core/pull/6938))
+  - `AccountsController:selectedAccountChange` is no longer part of the list of allowed events
+  - `NetworkController:getState` refetches balances
+  - `TransactionController:unapprovedTransactionAdded` refetches balances
+  - `TransactionController:unapprovedTransactionAdded'` refetches balances for the transaction from address and network
+  - `TransactionController:transactionConfirmed` refetches balances for the transaction from address and network
+- Add optional `fetchingEnabled` callback to `AccountTrackerController` constructor to stop it from fetching balances ([#6938](https://github.com/MetaMask/core/pull/6938))
+
 ## [85.0.0]
 
 ### Added

packages/assets-controllers/src/AccountTrackerController.test.ts

@@ -8,6 +8,7 @@ import {
   type MessengerEvents,
   type MockAnyNamespace,
 } from '@metamask/messenger';
+import type { NetworkConfiguration } from '@metamask/network-controller';
 import {
   type NetworkClientId,
   type NetworkClientConfiguration,
@@ -167,7 +168,7 @@ describe('AccountTrackerController', () => {
         mockedGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({
           tokenBalances: {
             '0x0000000000000000000000000000000000000000': {
-              [CHECKSUM_ADDRESS_1]: new BN('123456', 16),
+              [ADDRESS_1]: new BN('123456', 16),
             },
           },
           stakedBalances: {},
@@ -208,7 +209,7 @@ describe('AccountTrackerController', () => {
         mockedGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({
           tokenBalances: {
             '0x0000000000000000000000000000000000000000': {
-              [CHECKSUM_ADDRESS_1]: new BN('abcdef', 16),
+              [ADDRESS_1]: new BN('abcdef', 16),
             },
           },
           stakedBalances: {},
@@ -239,7 +240,88 @@ describe('AccountTrackerController', () => {
     );
   });
 
+  it('refreshes addresses when network is added', async () => {
+    await withController(
+      {
+        selectedAccount: ACCOUNT_1,
+        listAccounts: [ACCOUNT_1],
+      },
+      async ({ controller, messenger }) => {
+        mockedGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({
+          tokenBalances: {
+            '0x0000000000000000000000000000000000000000': {
+              [CHECKSUM_ADDRESS_1]: new BN('abcdef', 16),
+            },
+          },
+          stakedBalances: {},
+        });
+
+        messenger.publish(
+          'NetworkController:networkAdded',
+          {} as NetworkConfiguration,
+        );
+
+        await clock.tickAsync(1);
+
+        expect(
+          controller.state.accountsByChainId['0x1'][CHECKSUM_ADDRESS_1].balance,
+        ).toBe('0xabcdef');
+      },
+    );
+  });
+
+  it('refreshes addresses when keyring is unlocked', async () => {
+    await withController(
+      {
+        selectedAccount: ACCOUNT_1,
+        listAccounts: [ACCOUNT_1],
+      },
+      async ({ controller, messenger }) => {
+        mockedGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({
+          tokenBalances: {
+            '0x0000000000000000000000000000000000000000': {
+              [CHECKSUM_ADDRESS_1]: new BN('abcdef', 16),
+            },
+          },
+          stakedBalances: {},
+        });
+
+        messenger.publish('KeyringController:unlock');
+
+        await clock.tickAsync(1);
+
+        expect(
+          controller.state.accountsByChainId['0x1'][CHECKSUM_ADDRESS_1].balance,
+        ).toBe('0xabcdef');
+      },
+    );
+  });
+
   describe('refresh', () => {
+    it('does not refresh when fetching is disabled', async () => {
+      const expectedState = {
+        accountsByChainId: {
+          '0x1': {
+            [CHECKSUM_ADDRESS_1]: { balance: '0x0' },
+            [CHECKSUM_ADDRESS_2]: { balance: '0x0' },
+          },
+        },
+      };
+
+      await withController(
+        {
+          options: { fetchingEnabled: () => false },
+          selectedAccount: ACCOUNT_1,
+          listAccounts: [ACCOUNT_1, ACCOUNT_2],
+        },
+        async ({ controller, refresh }) => {
+          await refresh(clock, ['mainnet'], true);
+
+          expect(controller.state).toStrictEqual(expectedState);
+        },
+      );
+    });
+
     describe('without networkClientId', () => {
       it('should sync addresses', async () => {
         await withController(
@@ -1613,6 +1695,8 @@ async function withController<ReturnValue>(
       'AccountsController:selectedEvmAccountChange',
       'TransactionController:unapprovedTransactionAdded',
       'TransactionController:transactionConfirmed',
+      'NetworkController:networkAdded',
+      'KeyringController:unlock',
     ],
   });
 

packages/assets-controllers/src/AccountTrackerController.ts

@@ -3,7 +3,6 @@ import type {
   AccountsControllerSelectedEvmAccountChangeEvent,
   AccountsControllerGetSelectedAccountAction,
   AccountsControllerListAccountsAction,
-  AccountsControllerSelectedAccountChangeEvent,
 } from '@metamask/accounts-controller';
 import type {
   ControllerStateChangeEvent,
@@ -16,16 +15,17 @@ import {
   toChecksumHexAddress,
 } from '@metamask/controller-utils';
 import EthQuery from '@metamask/eth-query';
+import type { KeyringControllerUnlockEvent } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import type { Messenger } from '@metamask/messenger';
 import type {
   NetworkClient,
   NetworkClientId,
   NetworkControllerGetNetworkClientByIdAction,
   NetworkControllerGetStateAction,
+  NetworkControllerNetworkAddedEvent,
 } from '@metamask/network-controller';
 import { StaticIntervalPollingController } from '@metamask/polling-controller';
-import type { PreferencesControllerGetStateAction } from '@metamask/preferences-controller';
 import type {
   TransactionControllerTransactionConfirmedEvent,
   TransactionControllerUnapprovedTransactionAddedEvent,
@@ -174,7 +174,10 @@ export type AccountTrackerControllerActions =
  */
 export type AllowedActions =
   | AccountsControllerListAccountsAction
-  | PreferencesControllerGetStateAction
+  | {
+      type: 'PreferencesController:getState';
+      handler: () => { isMultiAccountBalancesEnabled: boolean };
+    }
   | AccountsControllerGetSelectedAccountAction
   | NetworkControllerGetStateAction
   | NetworkControllerGetNetworkClientByIdAction;
@@ -199,9 +202,10 @@ export type AccountTrackerControllerEvents =
  */
 export type AllowedEvents =
   | AccountsControllerSelectedEvmAccountChangeEvent
-  | AccountsControllerSelectedAccountChangeEvent
   | TransactionControllerUnapprovedTransactionAddedEvent
-  | TransactionControllerTransactionConfirmedEvent;
+  | TransactionControllerTransactionConfirmedEvent
+  | NetworkControllerNetworkAddedEvent
+  | KeyringControllerUnlockEvent;
 
 /**
  * The messenger of the {@link AccountTrackerController}.
@@ -236,6 +240,8 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
 
   readonly #balanceFetchers: BalanceFetcher[];
 
+  readonly #fetchingEnabled: () => boolean;
+
   /**
    * Creates an AccountTracker instance.
    *
@@ -247,6 +253,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
    * @param options.includeStakedAssets - Whether to include staked assets in the account balances.
    * @param options.accountsApiChainIds - Function that returns array of chainIds that should use Accounts-API strategy (if supported by API).
    * @param options.allowExternalServices - Disable external HTTP calls (privacy / offline mode).
+   * @param options.fetchingEnabled - Function that returns whether the controller is fetching enabled.
    */
   constructor({
     interval = 10000,
@@ -256,6 +263,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
     includeStakedAssets = false,
     accountsApiChainIds = () => [],
     allowExternalServices = () => true,
+    fetchingEnabled = () => true,
   }: {
     interval?: number;
     state?: Partial<AccountTrackerControllerState>;
@@ -264,6 +272,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
     includeStakedAssets?: boolean;
     accountsApiChainIds?: () => ChainIdHex[];
     allowExternalServices?: () => boolean;
+    fetchingEnabled?: () => boolean;
   }) {
     const { selectedNetworkClientId } = messenger.call(
       'NetworkController:getState',
@@ -302,6 +311,8 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
       ),
     ];
 
+    this.#fetchingEnabled = fetchingEnabled;
+
     this.setIntervalLength(interval);
 
     this.messenger.subscribe(
@@ -316,23 +327,39 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
       (event): string => event.address,
     );
 
+    this.messenger.subscribe('NetworkController:networkAdded', async () => {
+      await this.refresh(this.#getNetworkClientIds());
+    });
+
+    this.messenger.subscribe('KeyringController:unlock', async () => {
+      await this.refresh(this.#getNetworkClientIds());
+    });
+
     this.messenger.subscribe(
       'TransactionController:unapprovedTransactionAdded',
       async (transactionMeta: TransactionMeta) => {
-        await this.#refreshAddress(
-          [transactionMeta.networkClientId],
-          transactionMeta.txParams.from,
-        );
+        const addresses = [transactionMeta.txParams.from];
+        if (transactionMeta.txParams.to) {
+          addresses.push(transactionMeta.txParams.to);
+        }
+        await this.refreshAddresses({
+          networkClientIds: [transactionMeta.networkClientId],
+          addresses,
+        });
       },
     );
 
     this.messenger.subscribe(
       'TransactionController:transactionConfirmed',
       async (transactionMeta: TransactionMeta) => {
-        await this.#refreshAddress(
-          [transactionMeta.networkClientId],
-          transactionMeta.txParams.from,
-        );
+        const addresses = [transactionMeta.txParams.from];
+        if (transactionMeta.txParams.to) {
+          addresses.push(transactionMeta.txParams.to);
+        }
+        await this.refreshAddresses({
+          networkClientIds: [transactionMeta.networkClientId],
+          addresses,
+        });
       },
     );
 
@@ -540,13 +567,28 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
     });
   }
 
-  async #refreshAddress(networkClientIds: NetworkClientId[], address: string) {
-    const checksumAddress = toChecksumHexAddress(address) as ChecksumAddress;
+  async refreshAddresses({
+    networkClientIds,
+    addresses,
+  }: {
+    networkClientIds: NetworkClientId[];
+    addresses: string[];
+  }) {
+    const checksummedAddresses = addresses.map((address) =>
+      toChecksumHexAddress(address),
+    );
+
+    const accounts = this.messenger
+      .call('AccountsController:listAccounts')
+      .filter((account) =>
+        checksummedAddresses.includes(toChecksumHexAddress(account.address)),
+      );
+
     await this.#refreshAccounts({
       networkClientIds,
-      queryAllAccounts: false,
-      selectedAccount: checksumAddress,
-      allAccounts: [],
+      queryAllAccounts: true,
+      selectedAccount: '0x0',
+      allAccounts: accounts,
     });
   }
 
@@ -570,6 +612,10 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
 
       this.syncAccounts(chainIds);
 
+      if (!this.#fetchingEnabled()) {
+        return;
+      }
+
       // Use balance fetchers with fallback strategy
       const aggregated: ProcessedBalance[] = [];
       let remainingChains = [...chainIds] as ChainIdHex[];