fix: improve multichainTokenBalances and multichainAssetsRates controller state updates

packages/assets-controllers/CHANGELOG.md

@@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add event `MultichainAssetsController:accountAssetListUpdated` in MultichainAssetsController to notify when new assets are detected for an account ([#5761](https://github.com/MetaMask/core/pull/5761))
+
+### Changed
+
+- **BREAKING:** Removed subscription to `MultichainAssetsController:stateChange` in `MultichainAssetsRatesController` and add subscription to `MultichainAssetsController:accountAssetListUpdated` ([#5761](https://github.com/MetaMask/core/pull/5761))
+- **BREAKING:** Removed subscription to `MultichainAssetsController:stateChange` in `MultichainBalancesController` and add subscription to `MultichainAssetsController:accountAssetListUpdated` ([#5761](https://github.com/MetaMask/core/pull/5761))
+
 ## [61.1.0]
 
 ### Changed

packages/assets-controllers/src/MultichainAssetsController/MultichainAssetsController.ts

@@ -30,7 +30,6 @@ import type {
 import type { FungibleAssetMetadata, Snap, SnapId } from '@metamask/snaps-sdk';
 import { HandlerType } from '@metamask/snaps-utils';
 import {
-  hasProperty,
   isCaipAssetType,
   parseCaipAssetType,
   type CaipChainId,
@@ -57,6 +56,11 @@ export type AssetMetadataResponse = {
   };
 };
 
+export type MultichainAssetsControllerAccountAssetListUpdatedEvent = {
+  type: `${typeof controllerName}:accountAssetListUpdated`;
+  payload: AccountsControllerAccountAssetListUpdatedEvent['payload'];
+};
+
 /**
  * Constructs the default {@link MultichainAssetsController} state. This allows
  * consumers to provide a partial state object when initializing the controller
@@ -102,7 +106,8 @@ export type MultichainAssetsControllerActions =
  * Events emitted by {@link MultichainAssetsController}.
  */
 export type MultichainAssetsControllerEvents =
-  MultichainAssetsControllerStateChangeEvent;
+  | MultichainAssetsControllerStateChangeEvent
+  | MultichainAssetsControllerAccountAssetListUpdatedEvent;
 
 /**
  * A function executed within a mutually exclusive lock, with
@@ -254,32 +259,69 @@ export class MultichainAssetsController extends BaseController<
   ) {
     this.#assertControllerMutexIsLocked();
 
-    const assetsToUpdate = event.assets;
-    let assetsForMetadataRefresh = new Set<CaipAssetType>([]);
-    for (const accountId in assetsToUpdate) {
-      if (hasProperty(assetsToUpdate, accountId)) {
-        const { added, removed } = assetsToUpdate[accountId];
-        if (added.length > 0 || removed.length > 0) {
-          const existing = this.state.accountsAssets[accountId] || [];
-          const assets = new Set<CaipAssetType>([
-            ...existing,
-            ...added.filter((asset) => isCaipAssetType(asset)),
-          ]);
-          for (const removedAsset of removed) {
-            assets.delete(removedAsset);
-          }
-          assetsForMetadataRefresh = new Set([
-            ...assetsForMetadataRefresh,
-            ...assets,
-          ]);
-          this.update((state) => {
-            state.accountsAssets[accountId] = Array.from(assets);
-          });
+    const assetsForMetadataRefresh = new Set<CaipAssetType>([]);
+    const accountsAndAssetsToUpdate: AccountAssetListUpdatedEventPayload['assets'] =
+      {};
+    for (const [accountId, { added, removed }] of Object.entries(
+      event.assets,
+    )) {
+      if (added.length > 0 || removed.length > 0) {
+        const existing = this.state.accountsAssets[accountId] || [];
+
+        // In case accountsAndAssetsToUpdate event is fired with "added" assets that already exist, we don't want to add them again
+        const filteredToBeAddedAssets = added.filter(
+          (asset) => !existing.includes(asset) && isCaipAssetType(asset),
+        );
+
+        // In case accountsAndAssetsToUpdate event is fired with "removed" assets that don't exist, we don't want to remove them
+        const filteredToBeRemovedAssets = removed.filter(
+          (asset) => existing.includes(asset) && isCaipAssetType(asset),
+        );
+
+        if (
+          filteredToBeAddedAssets.length > 0 ||
+          filteredToBeRemovedAssets.length > 0
+        ) {
+          accountsAndAssetsToUpdate[accountId] = {
+            added: filteredToBeAddedAssets,
+            removed: filteredToBeRemovedAssets,
+          };
+        }
+
+        for (const asset of existing) {
+          assetsForMetadataRefresh.add(asset);
+        }
+        for (const asset of filteredToBeAddedAssets) {
+          assetsForMetadataRefresh.add(asset);
+        }
+        for (const asset of filteredToBeRemovedAssets) {
+          assetsForMetadataRefresh.delete(asset);
         }
       }
     }
+
+    this.update((state) => {
+      for (const [accountId, { added, removed }] of Object.entries(
+        accountsAndAssetsToUpdate,
+      )) {
+        const assets = new Set([
+          ...(state.accountsAssets[accountId] || []),
+          ...added,
+        ]);
+        for (const asset of removed) {
+          assets.delete(asset);
+        }
+
+        state.accountsAssets[accountId] = Array.from(assets);
+      }
+    });
+
     // Trigger fetching metadata for new assets
     await this.#refreshAssetsMetadata(Array.from(assetsForMetadataRefresh));
+
+    this.messagingSystem.publish(`${controllerName}:accountAssetListUpdated`, {
+      assets: accountsAndAssetsToUpdate,
+    });
   }
 
   /**
@@ -318,6 +360,17 @@ export class MultichainAssetsController extends BaseController<
       this.update((state) => {
         state.accountsAssets[account.id] = assets;
       });
+      this.messagingSystem.publish(
+        `${controllerName}:accountAssetListUpdated`,
+        {
+          assets: {
+            [account.id]: {
+              added: assets,
+              removed: [],
+            },
+          },
+        },
+      );
     }
   }
 

packages/assets-controllers/src/MultichainAssetsController/index.ts

@@ -8,6 +8,7 @@ export type {
   MultichainAssetsControllerGetStateAction,
   MultichainAssetsControllerStateChangeEvent,
   MultichainAssetsControllerActions,
-  MultichainAssetsControllerEvents,
   MultichainAssetsControllerMessenger,
+  MultichainAssetsControllerAccountAssetListUpdatedEvent,
+  MultichainAssetsControllerEvents,
 } from './MultichainAssetsController';

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

@@ -1,14 +1,20 @@
 import { Messenger } from '@metamask/base-controller';
+import { SolScope } from '@metamask/keyring-api';
+import { SolMethod } from '@metamask/keyring-api';
+import { SolAccountType } from '@metamask/keyring-api';
+import { KeyringTypes } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import { KeyringClient } from '@metamask/keyring-snap-client';
 import type { OnAssetHistoricalPriceResponse } from '@metamask/snaps-sdk';
 import { useFakeTimers } from 'sinon';
+import { v4 as uuidv4 } from 'uuid';
 
 import { MultichainAssetsRatesController } from '.';
 import {
   type AllowedActions,
   type AllowedEvents,
 } from './MultichainAssetsRatesController';
+import { advanceTime } from '../../../../tests/helpers';
 
 // A fake non‑EVM account (with Snap metadata) that meets the controller’s criteria.
 const fakeNonEvmAccount: InternalAccount = {
@@ -163,16 +169,21 @@ const setupController = ({
       'KeyringController:lock',
       'KeyringController:unlock',
       'CurrencyRateController:stateChange',
-      'MultichainAssetsController:stateChange',
+      'MultichainAssetsController:accountAssetListUpdated',
     ],
   });
 
+  const controller = new MultichainAssetsRatesController({
+    messenger: multichainAssetsRatesControllerMessenger,
+    ...config,
+  });
+
+  const updateSpy = jest.spyOn(controller, 'update' as never);
+
   return {
-    controller: new MultichainAssetsRatesController({
-      messenger: multichainAssetsRatesControllerMessenger,
-      ...config,
-    }),
+    controller,
     messenger,
+    updateSpy,
   };
 };
 
@@ -282,24 +293,6 @@ describe('MultichainAssetsRatesController', () => {
     expect(snapHandler).not.toHaveBeenCalled();
   });
 
-  it('does not update conversion rates if the assets are empty', async () => {
-    const { controller, messenger } = setupController();
-
-    const snapSpy = jest.fn().mockResolvedValue({ conversionRates: {} });
-    messenger.registerActionHandler('SnapController:handleRequest', snapSpy);
-
-    // Publish a selectedAccountChange event.
-    // @ts-expect-error-next-line
-    messenger.publish('MultichainAssetsController:stateChange', {
-      accountsAssets: {
-        account3: [],
-      },
-    });
-
-    expect(snapSpy).not.toHaveBeenCalled();
-    expect(controller.state.conversionRates).toStrictEqual({});
-  });
-
   it('resumes update tokens rates when the keyring is unlocked', async () => {
     const { controller, messenger } = setupController();
     messenger.publish('KeyringController:lock');
@@ -352,24 +345,108 @@ describe('MultichainAssetsRatesController', () => {
     expect(updateSpy).toHaveBeenCalled();
   });
 
-  it('calls updateTokensRates when an multichain assets state is updated', async () => {
-    const { controller, messenger } = setupController();
+  it('calls updateTokensRatesForNewAssets when newAccountAssets event is published', async () => {
+    const testAccounts = [
+      {
+        address: 'EBBYfhQzVzurZiweJ2keeBWpgGLs1cbWYcz28gjGgi5x',
+        id: uuidv4(),
+        metadata: {
+          name: 'Solana Account 1',
+          importTime: Date.now(),
+          keyring: {
+            type: KeyringTypes.snap,
+          },
+          snap: {
+            id: 'mock-sol-snap',
+            name: 'mock-sol-snap',
+            enabled: true,
+          },
+          lastSelected: 0,
+        },
+        scopes: [SolScope.Devnet],
+        options: {},
+        methods: [SolMethod.SendAndConfirmTransaction],
+        type: SolAccountType.DataAccount,
+      },
+      {
+        address: 'GMTYfhQzVzurZiweJ2keeBWpgGLs1cbWYcz28gjGgi5x',
+        id: uuidv4(),
+        metadata: {
+          name: 'Solana Account 2',
+          importTime: Date.now(),
+          keyring: {
+            type: KeyringTypes.snap,
+          },
+          snap: {
+            id: 'mock-sol-snap',
+            name: 'mock-sol-snap',
+            enabled: true,
+          },
+          lastSelected: 0,
+        },
+        scopes: [SolScope.Devnet],
+        options: {},
+        methods: [SolMethod.SendAndConfirmTransaction],
+        type: SolAccountType.DataAccount,
+      },
+    ];
+    const { controller, messenger, updateSpy } = setupController({
+      accountsAssets: testAccounts,
+    });
 
-    // Spy on updateTokensRates.
-    const updateSpy = jest
-      .spyOn(controller, 'updateAssetsRates')
-      .mockResolvedValue();
+    const snapSpy = jest
+      .fn()
+      .mockResolvedValueOnce({
+        conversionRates: {
+          'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501': {
+            'swift:0/iso4217:USD': {
+              rate: '100',
+              conversionTime: 1738539923277,
+            },
+          },
+        },
+      })
+      .mockResolvedValueOnce({
+        conversionRates: {
+          'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token1:501': {
+            'swift:0/iso4217:USD': {
+              rate: '200',
+              conversionTime: 1738539923277,
+            },
+          },
+        },
+      });
+    messenger.registerActionHandler('SnapController:handleRequest', snapSpy);
 
-    // Publish a selectedAccountChange event.
-    // @ts-expect-error-next-line
-    messenger.publish('MultichainAssetsController:stateChange', {
-      accountsAssets: {
-        account3: ['solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501'],
+    messenger.publish('MultichainAssetsController:accountAssetListUpdated', {
+      assets: {
+        [testAccounts[0].id]: {
+          added: ['solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501'],
+          removed: [],
+        },
+        [testAccounts[1].id]: {
+          added: ['solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token1:501'],
+          removed: [],
+        },
       },
     });
     // Wait for the asynchronous subscriber to run.
     await Promise.resolve();
-    expect(updateSpy).toHaveBeenCalled();
+    await advanceTime({ clock, duration: 10 });
+
+    expect(updateSpy).toHaveBeenCalledTimes(1);
+    expect(controller.state.conversionRates).toMatchObject({
+      'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501': {
+        rate: '100',
+        conversionTime: 1738539923277,
+        currency: 'swift:0/iso4217:USD',
+      },
+      'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token1:501': {
+        rate: '200',
+        conversionTime: 1738539923277,
+        currency: 'swift:0/iso4217:USD',
+      },
+    });
   });
 
   it('handles partial or empty Snap responses gracefully', async () => {
@@ -436,28 +513,6 @@ describe('MultichainAssetsRatesController', () => {
     expect(updateSpy).toHaveBeenCalled();
   });
 
-  it('should return an empty array if no assets are found', async () => {
-    const { controller, messenger } = setupController();
-
-    const snapSpy = jest.fn().mockResolvedValue({ conversionRates: {} });
-    messenger.registerActionHandler('SnapController:handleRequest', snapSpy);
-
-    messenger.publish(
-      'MultichainAssetsController:stateChange',
-      {
-        accountsAssets: {
-          account1: [],
-        },
-        assetsMetadata: {},
-      },
-      [],
-    );
-
-    await controller.updateAssetsRates();
-
-    expect(controller.state.conversionRates).toStrictEqual({});
-  });
-
   describe('fetchHistoricalPricesForAsset', () => {
     it('throws an error if call to snap fails', async () => {
       const testAsset = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501';