feat: assets controllers package upgrade with fix

.storybook/test-data.js

@@ -792,47 +792,26 @@ const state = {
     swapsWelcomeMessageHasBeenShown: true,
     defaultHomeActiveTabName: 'Tokens',
     network: '5',
-    accounts: {
-      '0x64a845a5b02460acf8a3d84503b0d68d028b4bb4': {
-        address: '0x64a845a5b02460acf8a3d84503b0d68d028b4bb4',
-        balance: '0x176e5b6f173ebe66',
-      },
-      '0xb19ac54efa18cc3a14a5b821bfec73d284bf0c5e': {
-        address: '0xb19ac54efa18cc3a14a5b821bfec73d284bf0c5e',
-        balance: '0x2d3142f5000',
-      },
-      '0x9d0ba4ddac06032527b140912ec808ab9451b788': {
-        address: '0x9d0ba4ddac06032527b140912ec808ab9451b788',
-        balance: '0x15f6f0b9d4f8d000',
-      },
-    },
     accountsByChainId: {
       '0x1': {
-        '0x64a845a5b02460acf8a3d84503b0d68d028b4bb4': { balance: '0x0' },
-        '0xb19ac54efa18cc3a14a5b821bfec73d284bf0c5e': {
+        '0x64A845a5b02460ACf8a3D84503b0D68d028B4bb4': { balance: '0x0' },
+        '0xb19Ac54EfA18CC3A14A5B821bFeC73d284Bf0c5e': {
           balance: '0xcaf5317161f400',
         },
-        '0x9d0ba4ddac06032527b140912ec808ab9451b788': { balance: '0x0' },
+        '0x9D0ba4DDAC06032527B140912EC808ab9451b788': { balance: '0x0' },
       },
       '0x5': {
-        '0x64a845a5b02460acf8a3d84503b0d68d028b4bb4': {
-          address: '0x64a845a5b02460acf8a3d84503b0d68d028b4bb4',
+        '0x64A845a5b02460ACf8a3D84503b0D68d028B4bb4': {
           balance: '0x176e5b6f173ebe66',
         },
-        '0xb19ac54efa18cc3a14a5b821bfec73d284bf0c5e': {
-          address: '0xb19ac54efa18cc3a14a5b821bfec73d284bf0c5e',
+        '0xb19Ac54EfA18CC3A14A5B821bFeC73d284Bf0c5e': {
           balance: '0x2d3142f5000',
         },
-        '0x9d0ba4ddac06032527b140912ec808ab9451b788': {
-          address: '0x9d0ba4ddac06032527b140912ec808ab9451b788',
+        '0x9D0ba4DDAC06032527B140912EC808ab9451b788': {
           balance: '0x15f6f0b9d4f8d000',
         },
       },
     },
-    currentBlockGasLimit: '0x793af4',
-    currentBlockGasLimitByChainId: {
-      '0x5': '0x793af4',
-    },
     transactions: [
       {
         chainId: '0x38',

app/scripts/background.js

@@ -556,10 +556,7 @@ function saveTimestamp() {
  * @property {boolean} welcomeScreen - True if welcome screen should be shown.
  * @property {string} currentLocale - A locale string matching the user's preferred display language.
  * @property {string} networkStatus - Either "unknown", "available", "unavailable", or "blocked", depending on the status of the currently selected network.
- * @property {object} accounts - An object mapping lower-case hex addresses to objects with "balance" and "address" keys, both storing hex string values.
  * @property {object} accountsByChainId - An object mapping lower-case hex addresses to objects with "balance" and "address" keys, both storing hex string values keyed by chain id.
- * @property {hex} currentBlockGasLimit - The most recently seen block gas limit, in a lower case hex prefixed string.
- * @property {object} currentBlockGasLimitByChainId - The most recently seen block gas limit, in a lower case hex prefixed string keyed by chain id.
  * @property {object} unapprovedPersonalMsgs - An object of messages pending approval, mapping a unique ID to the options.
  * @property {number} unapprovedPersonalMsgCount - The number of messages in unapprovedPersonalMsgs.
  * @property {object} unapprovedEncryptionPublicKeyMsgs - An object of messages pending approval, mapping a unique ID to the options.

app/scripts/constants/sentry-state.ts

@@ -14,10 +14,7 @@ export const SENTRY_BACKGROUND_STATE = {
     },
   },
   AccountTracker: {
-    accounts: false,
     accountsByChainId: false,
-    currentBlockGasLimit: true,
-    currentBlockGasLimitByChainId: true,
   },
   AddressBookController: {
     addressBook: false,

app/scripts/controller-init/account-tracker-controller-init.test.ts

@@ -1,10 +1,6 @@
 import { ActionConstraint, Messenger } from '@metamask/base-controller';
-import {
-  NetworkControllerGetSelectedNetworkClientAction,
-  NetworkControllerNetworkDidChangeEvent,
-} from '@metamask/network-controller';
 import { RemoteFeatureFlagControllerGetStateAction } from '@metamask/remote-feature-flag-controller';
-import AccountTrackerController from '../controllers/account-tracker-controller';
+import { AccountTrackerController } from '@metamask/assets-controllers';
 import { ControllerInitRequest } from './types';
 import { buildControllerInitRequestMock } from './test/utils';
 import {
@@ -19,9 +15,7 @@ jest.mock('../controllers/account-tracker-controller');
 
 function getInitRequestMock(
   baseMessenger = new Messenger<
-    | NetworkControllerGetSelectedNetworkClientAction
-    | RemoteFeatureFlagControllerGetStateAction
-    | ActionConstraint,
+    RemoteFeatureFlagControllerGetStateAction | ActionConstraint,
     never
   >(),
 ): jest.Mocked<
@@ -30,17 +24,6 @@ function getInitRequestMock(
     AccountTrackerControllerInitMessenger
   >
 > {
-  baseMessenger.registerActionHandler(
-    'NetworkController:getSelectedNetworkClient',
-    () => ({
-      // @ts-expect-error: Partial mock.
-      provider: {},
-
-      // @ts-expect-error: Partial mock.
-      blockTracker: {},
-    }),
-  );
-
   baseMessenger.registerActionHandler(
     'RemoteFeatureFlagController:getState',
     () => ({
@@ -93,26 +76,4 @@ describe('AccountTrackerControllerInit', () => {
     expect(chainIds).toContain('0x38'); // BSC
     expect(chainIds).toContain('0xe708'); // Linea
   });
-
-  it('calls `updateAccounts` when `NetworkController:networkDidChange` is emitted', () => {
-    const messenger = new Messenger<
-      | NetworkControllerGetSelectedNetworkClientAction
-      | RemoteFeatureFlagControllerGetStateAction
-      | ActionConstraint,
-      NetworkControllerNetworkDidChangeEvent
-    >();
-
-    const request = getInitRequestMock(messenger);
-    const { controller } = AccountTrackerControllerInit(request);
-
-    expect(controller.updateAccounts).not.toHaveBeenCalled();
-
-    messenger.publish('NetworkController:networkDidChange', {
-      selectedNetworkClientId: 'test',
-      networkConfigurationsByChainId: {},
-      networksMetadata: {},
-    });
-
-    expect(controller.updateAccounts).toHaveBeenCalledTimes(1);
-  });
 });

app/scripts/controller-init/account-tracker-controller-init.ts

@@ -1,56 +1,38 @@
-import { assert } from '@metamask/utils';
-import { getProviderConfig } from '../../../shared/modules/selectors/networks';
-import { NETWORK_TYPES } from '../../../shared/constants/network';
-import AccountTrackerController from '../controllers/account-tracker-controller';
+import { AccountTrackerController } from '@metamask/assets-controllers';
+import { NetworkClientId } from '@metamask/network-controller';
 import { ControllerInitFunction } from './types';
 import {
   AccountTrackerControllerInitMessenger,
   AccountTrackerControllerMessenger,
 } from './messengers';
 
-/**
- * Initialize the account tracker controller.
- *
- * @param request - The request object.
- * @param request.controllerMessenger - The messenger to use for the controller.
- * @param request.initMessenger - The messenger to use for initialization.
- * @returns The initialized controller.
- */
 export const AccountTrackerControllerInit: ControllerInitFunction<
   AccountTrackerController,
   AccountTrackerControllerMessenger,
   AccountTrackerControllerInitMessenger
-> = ({ controllerMessenger, initMessenger }) => {
-  const { provider, blockTracker } =
-    initMessenger.call('NetworkController:getSelectedNetworkClient') ?? {};
-
-  assert(
-    provider,
-    'Provider is required to initialize AccountTrackerController.',
-  );
-
-  assert(
-    blockTracker,
-    'Block tracker is required to initialize AccountTrackerController.',
-  );
+> = ({ controllerMessenger, initMessenger, persistedState, getController }) => {
+  const getAssetsContractController = () =>
+    getController('AssetsContractController');
 
   const controller = new AccountTrackerController({
-    state: { accounts: {} },
+    state: persistedState.AccountTrackerController,
     messenger: controllerMessenger,
-    provider,
-    blockTracker,
-    getNetworkIdentifier: (providerConfig): string => {
-      const metamask = initMessenger.call('NetworkController:getState');
-
-      const config =
-        providerConfig ??
-        getProviderConfig({
-          metamask,
-        });
-
-      return config.type === NETWORK_TYPES.RPC && config.rpcUrl
-        ? config.rpcUrl
-        : config.type;
+    getStakedBalanceForChain: (
+      addresses: string[],
+      networkClientId?: NetworkClientId,
+    ) => {
+      const assetsContractController = getAssetsContractController();
+      return assetsContractController.getStakedBalanceForChain(
+        addresses,
+        networkClientId,
+      );
+    },
+    includeStakedAssets: false,
+    allowExternalServices: () => {
+      const { useExternalServices } = initMessenger.call(
+        'PreferencesController:getState',
+      );
+      return useExternalServices;
     },
     accountsApiChainIds: () => {
       const state = initMessenger.call('RemoteFeatureFlagController:getState');
@@ -59,16 +41,11 @@ export const AccountTrackerControllerInit: ControllerInitFunction<
         state?.remoteFeatureFlags?.assetsAccountApiBalances;
 
       return Array.isArray(featureFlagForAccountApiBalances)
-        ? (featureFlagForAccountApiBalances as string[])
+        ? (featureFlagForAccountApiBalances as `0x${string}`[])
         : [];
     },
   });
 
-  // Ensure `AccountTrackerController` updates balances after network change.
-  initMessenger.subscribe('NetworkController:networkDidChange', () => {
-    controller.updateAccounts();
-  });
-
   return {
     persistedStateKey: null,
     memStateKey: null,

app/scripts/controller-init/confirmations/transaction-controller-init.ts

@@ -56,7 +56,6 @@ export const TransactionControllerInit: ControllerInitFunction<
     getFlatState,
     getPermittedAccounts,
     getTransactionMetricsRequest,
-    updateAccountBalanceForTransactionNetwork,
     persistedState,
   } = request;
 
@@ -191,7 +190,6 @@ export const TransactionControllerInit: ControllerInitFunction<
   addTransactionControllerListeners(
     initMessenger,
     getTransactionMetricsRequest,
-    updateAccountBalanceForTransactionNetwork,
   );
 
   const api = getApi(controller);
@@ -259,22 +257,9 @@ function getExternalPendingTransactions(
 function addTransactionControllerListeners(
   initMessenger: TransactionControllerInitMessenger,
   getTransactionMetricsRequest: () => TransactionMetricsRequest,
-  updateAccountBalanceForTransactionNetwork: (
-    transactionMeta: TransactionMeta,
-  ) => void,
 ) {
   const transactionMetricsRequest = getTransactionMetricsRequest();
 
-  initMessenger.subscribe(
-    'TransactionController:unapprovedTransactionAdded',
-    updateAccountBalanceForTransactionNetwork,
-  );
-
-  initMessenger.subscribe(
-    'TransactionController:transactionConfirmed',
-    updateAccountBalanceForTransactionNetwork,
-  );
-
   initMessenger.subscribe(
     'TransactionController:postTransactionBalanceUpdated',
     // TODO: Fix in https://github.com/MetaMask/metamask-extension/issues/31879

app/scripts/controller-init/controller-list.ts

@@ -13,6 +13,7 @@ import { SmartTransactionsController } from '@metamask/smart-transactions-contro
 import { TransactionController } from '@metamask/transaction-controller';
 import { AccountsController } from '@metamask/accounts-controller';
 import {
+  AccountTrackerController,
   AssetsContractController,
   CurrencyRateController,
   DeFiPositionsController,
@@ -88,7 +89,6 @@ import { NetworkOrderController } from '../controllers/network-order';
 import OAuthService from '../services/oauth/oauth-service';
 import MetaMetricsController from '../controllers/metametrics-controller';
 import { SnapsNameProvider } from '../lib/SnapsNameProvider';
-import AccountTrackerController from '../controllers/account-tracker-controller';
 import { AppStateController } from '../controllers/app-state-controller';
 import { SnapKeyringBuilder } from '../lib/snap-keyring/snap-keyring';
 import { SubscriptionService } from '../services/subscription/subscription-service';
@@ -257,4 +257,5 @@ export type ControllerFlatState = AccountOrderController['state'] &
   TokenRatesController['state'] &
   NftController['state'] &
   NftDetectionController['state'] &
-  NetworkEnablementController['state'];
+  NetworkEnablementController['state'] &
+  AccountTrackerController['state'];