diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index 331d421f9ba..669d18932f1 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/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 diff --git a/packages/assets-controllers/src/MultichainAssetsController/MultichainAssetsController.ts b/packages/assets-controllers/src/MultichainAssetsController/MultichainAssetsController.ts index bc40d803ae5..fc4cbaa49b3 100644 --- a/packages/assets-controllers/src/MultichainAssetsController/MultichainAssetsController.ts +++ b/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([]); - 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([ - ...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([]); + 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: [], + }, + }, + }, + ); } } diff --git a/packages/assets-controllers/src/MultichainAssetsController/index.ts b/packages/assets-controllers/src/MultichainAssetsController/index.ts index a558a58720d..bfe10978eb8 100644 --- a/packages/assets-controllers/src/MultichainAssetsController/index.ts +++ b/packages/assets-controllers/src/MultichainAssetsController/index.ts @@ -8,6 +8,7 @@ export type { MultichainAssetsControllerGetStateAction, MultichainAssetsControllerStateChangeEvent, MultichainAssetsControllerActions, - MultichainAssetsControllerEvents, MultichainAssetsControllerMessenger, + MultichainAssetsControllerAccountAssetListUpdatedEvent, + MultichainAssetsControllerEvents, } from './MultichainAssetsController'; diff --git a/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.test.ts b/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.test.ts index 1da5feef59d..97ab0158970 100644 --- a/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.test.ts +++ b/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'; diff --git a/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.ts b/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.ts index e8e98d13774..d42c41acbf5 100644 --- a/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.ts +++ b/packages/assets-controllers/src/MultichainAssetsRatesController/MultichainAssetsRatesController.ts @@ -37,8 +37,8 @@ import type { } from '../CurrencyRateController'; import type { MultichainAssetsControllerGetStateAction, + MultichainAssetsControllerAccountAssetListUpdatedEvent, MultichainAssetsControllerState, - MultichainAssetsControllerStateChangeEvent, } from '../MultichainAssetsController'; /** @@ -132,8 +132,7 @@ export type AllowedEvents = | KeyringControllerUnlockEvent | AccountsControllerAccountAddedEvent | CurrencyRateStateChange - | MultichainAssetsControllerStateChangeEvent; - + | MultichainAssetsControllerAccountAssetListUpdatedEvent; /** * Messenger type for the MultichainAssetsRatesController. */ @@ -171,7 +170,7 @@ export class MultichainAssetsRatesController extends StaticIntervalPollingContro #currentCurrency: CurrencyRateState['currentCurrency']; - #accountsAssets: MultichainAssetsControllerState['accountsAssets']; + readonly #accountsAssets: MultichainAssetsControllerState['accountsAssets']; #isUnlocked = true; @@ -229,10 +228,16 @@ export class MultichainAssetsRatesController extends StaticIntervalPollingContro ); this.messagingSystem.subscribe( - 'MultichainAssetsController:stateChange', - async (multichainAssetsState: MultichainAssetsControllerState) => { - this.#accountsAssets = multichainAssetsState.accountsAssets; - await this.updateAssetsRates(); + 'MultichainAssetsController:accountAssetListUpdated', + async ({ assets }) => { + const newAccountAssets = Object.entries(assets).map( + ([accountId, { added }]) => ({ + accountId, + assets: [...added], + }), + ); + // TODO; removed can be used in future for further cleanup + await this.#updateAssetsRatesForNewAssets(newAccountAssets); }, ); } @@ -309,33 +314,41 @@ export class MultichainAssetsRatesController extends StaticIntervalPollingContro continue; } - // Build the conversions array - const conversions = this.#buildConversions(assets); - - // Retrieve rates from Snap - const accountRates: OnAssetsConversionResponse = - (await this.#handleSnapRequest({ - snapId: account?.metadata.snap?.id as SnapId, - handler: HandlerType.OnAssetsConversion, - params: { - ...conversions, - includeMarketData: true, - }, - })) as OnAssetsConversionResponse; - - // Flatten nested rates if needed - const flattenedRates = this.#flattenRates(accountRates); - - // Build the updatedRates object for these assets - const updatedRates = this.#buildUpdatedRates(assets, flattenedRates); + const rates = await this.#getUpdatedRatesFor(account, assets); // Apply these updated rates to controller state - this.#applyUpdatedRates(updatedRates); + this.#applyUpdatedRates(rates); } })().finally(() => { releaseLock(); }); } + async #getUpdatedRatesFor( + account: InternalAccount, + assets: CaipAssetType[], + ): Promise> { + // Build the conversions array + const conversions = this.#buildConversions(assets); + + // Retrieve rates from Snap + const accountRates: OnAssetsConversionResponse = + (await this.#handleSnapRequest({ + snapId: account?.metadata.snap?.id as SnapId, + handler: HandlerType.OnAssetsConversion, + params: { + ...conversions, + includeMarketData: true, + }, + })) as OnAssetsConversionResponse; + + // Flatten nested rates if needed + const flattenedRates = this.#flattenRates(accountRates); + + // Build the updatedRates object for these assets + const updatedRates = this.#buildUpdatedRates(assets, flattenedRates); + return updatedRates; + } + /** * Fetches historical prices for the current account * @@ -399,6 +412,63 @@ export class MultichainAssetsRatesController extends StaticIntervalPollingContro }); } + /** + * Updates the conversion rates for new assets. + * + * @param accounts - The accounts to update the conversion rates for. + * @returns A promise that resolves when the rates are updated. + */ + async #updateAssetsRatesForNewAssets( + accounts: { + accountId: string; + assets: CaipAssetType[]; + }[], + ): Promise { + const releaseLock = await this.#mutex.acquire(); + + return (async () => { + if (!this.isActive) { + return; + } + const allNewRates: Record< + string, + { rate: string | null; conversionTime: number | null } + > = {}; + + for (const { accountId, assets } of accounts) { + const account = this.#getAccount(accountId); + + const rates = await this.#getUpdatedRatesFor(account, assets); + // Track new rates + for (const [asset, rate] of Object.entries(rates)) { + allNewRates[asset] = rate; + } + } + + this.#applyUpdatedRates(allNewRates); + })().finally(() => { + releaseLock(); + }); + } + + /** + * Get a non-EVM account from its ID. + * + * @param accountId - The account ID. + * @returns The non-EVM account. + */ + #getAccount(accountId: string): InternalAccount { + const account: InternalAccount | undefined = this.#listAccounts().find( + (multichainAccount) => multichainAccount.id === accountId, + ); + + if (!account) { + throw new Error(`Unknown account: ${accountId}`); + } + + return account; + } + /** * Returns the array of CAIP-19 assets for the given account ID. * If none are found, returns an empty array. @@ -488,6 +558,9 @@ export class MultichainAssetsRatesController extends StaticIntervalPollingContro { rate: string | null; conversionTime: number | null } >, ): void { + if (Object.keys(updatedRates).length === 0) { + return; + } this.update((state: Draft) => { state.conversionRates = { ...state.conversionRates, diff --git a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.test.ts b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.test.ts index 8d92182256e..7fe8810d168 100644 --- a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.test.ts +++ b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.test.ts @@ -141,7 +141,7 @@ function getRestrictedMessenger( 'AccountsController:accountAdded', 'AccountsController:accountRemoved', 'AccountsController:accountBalancesUpdated', - 'MultichainAssetsController:stateChange', + 'MultichainAssetsController:accountAssetListUpdated', ], }); } @@ -154,6 +154,11 @@ const setupController = ({ mocks?: { listMultichainAccounts?: InternalAccount[]; handleRequestReturnValue?: Record; + handleMockGetAssetsState?: { + accountsAssets: { + [account: string]: CaipAssetType[]; + }; + }; }; } = {}) => { const messenger = getRootMessenger(); @@ -175,11 +180,13 @@ const setupController = ({ ), ); - const mockGetAssetsState = jest.fn().mockReturnValue({ - accountsAssets: { - [mockBtcAccount.id]: [mockBtcNativeAsset], + const mockGetAssetsState = jest.fn().mockReturnValue( + mocks?.handleMockGetAssetsState ?? { + accountsAssets: { + [mockBtcAccount.id]: [mockBtcNativeAsset], + }, }, - }); + ); messenger.registerActionHandler( 'MultichainAssetsController:getState', mockGetAssetsState, @@ -221,7 +228,7 @@ async function waitForAllPromises(): Promise { await new Promise(process.nextTick); } -describe('BalancesController', () => { +describe('MultichainBalancesController', () => { it('initialize with default state', () => { const messenger = getRootMessenger(); const multichainBalancesMessenger = getRestrictedMessenger(messenger); @@ -419,25 +426,205 @@ describe('BalancesController', () => { expect(controller.state.balances[mockBtcAccount.id]).toStrictEqual({}); }); - it('updates balances when receiving "MultichainAssetsController:stateChange" event', async () => { - const { controller, messenger } = setupController(); - - messenger.publish( - 'MultichainAssetsController:stateChange', + describe('when "MultichainAssetsController:accountAssetListUpdated" is fired', () => { + const mockListSolanaAccounts = [ { - assetsMetadata: {}, - accountsAssets: { - [mockBtcAccount.id]: [mockBtcNativeAsset], + 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, + }, + ]; - await waitForAllPromises(); + it('updates balances when receiving "MultichainAssetsController:accountAssetListUpdated" event and state is empty', async () => { + const mockSolanaAccountId1 = mockListSolanaAccounts[0].id; + const mockSolanaAccountId2 = mockListSolanaAccounts[1].id; - expect(controller.state.balances[mockBtcAccount.id]).toStrictEqual( - mockBalanceResult, - ); + const { controller, messenger, mockSnapHandleRequest } = setupController({ + state: { + balances: {}, + }, + mocks: { + handleMockGetAssetsState: { + accountsAssets: {}, + }, + handleRequestReturnValue: {}, + listMultichainAccounts: mockListSolanaAccounts, + }, + }); + + mockSnapHandleRequest.mockReset(); + mockSnapHandleRequest + .mockResolvedValueOnce({ + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken': { + amount: '1.00000000', + unit: 'SOL', + }, + }) + .mockResolvedValueOnce({ + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3': { + amount: '3.00000000', + unit: 'SOL', + }, + }); + messenger.publish('MultichainAssetsController:accountAssetListUpdated', { + assets: { + [mockSolanaAccountId1]: { + added: ['solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken'], + removed: [], + }, + [mockSolanaAccountId2]: { + added: ['solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3'], + removed: [] as `${string}:${string}/${string}:${string}`[], + }, + }, + }); + + await waitForAllPromises(); + + expect(controller.state.balances).toStrictEqual({ + [mockSolanaAccountId1]: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken': { + amount: '1.00000000', + unit: 'SOL', + }, + }, + [mockSolanaAccountId2]: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3': { + amount: '3.00000000', + unit: 'SOL', + }, + }, + }); + }); + + it('updates balances when receiving "MultichainAssetsController:accountAssetListUpdated" event and state has existing balances', async () => { + const mockSolanaAccountId1 = mockListSolanaAccounts[0].id; + const mockSolanaAccountId2 = mockListSolanaAccounts[1].id; + + const existingBalancesState = { + [mockSolanaAccountId1]: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken55': { + amount: '5.00000000', + unit: 'SOL', + }, + }, + }; + const { + controller, + messenger, + mockSnapHandleRequest, + mockListMultichainAccounts, + } = setupController({ + state: { + balances: existingBalancesState, + }, + mocks: { + handleMockGetAssetsState: { + accountsAssets: { + [mockSolanaAccountId1]: [ + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken55', + ], + }, + }, + handleRequestReturnValue: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken55': { + amount: '55.00000000', + unit: 'SOL', + }, + }, + listMultichainAccounts: [mockListSolanaAccounts[0]], + }, + }); + + mockSnapHandleRequest.mockReset(); + mockListMultichainAccounts.mockReset(); + + mockListMultichainAccounts.mockReturnValue(mockListSolanaAccounts); + mockSnapHandleRequest + .mockResolvedValueOnce({ + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken': { + amount: '1.00000000', + unit: 'SOL', + }, + }) + .mockResolvedValueOnce({ + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3': { + amount: '3.00000000', + unit: 'SOL', + }, + }); + + messenger.publish('MultichainAssetsController:accountAssetListUpdated', { + assets: { + [mockSolanaAccountId1]: { + added: ['solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken'], + removed: [], + }, + [mockSolanaAccountId2]: { + added: ['solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3'], + removed: [], + }, + }, + }); + + await waitForAllPromises(); + + expect(controller.state.balances).toStrictEqual({ + [mockSolanaAccountId1]: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken': { + amount: '1.00000000', + unit: 'SOL', + }, + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken55': { + amount: '55.00000000', + unit: 'SOL', + }, + }, + [mockSolanaAccountId2]: { + 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1/token:newToken3': { + amount: '3.00000000', + unit: 'SOL', + }, + }, + }); + }); }); it('resumes updating balances after unlocking KeyringController', async () => { diff --git a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts index 18d523f0100..e39372874d1 100644 --- a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts +++ b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts @@ -27,8 +27,7 @@ import type { Draft } from 'immer'; import type { MultichainAssetsControllerGetStateAction, - MultichainAssetsControllerState, - MultichainAssetsControllerStateChangeEvent, + MultichainAssetsControllerAccountAssetListUpdatedEvent, } from '../MultichainAssetsController'; const controllerName = 'MultichainBalancesController'; @@ -105,8 +104,7 @@ type AllowedEvents = | AccountsControllerAccountAddedEvent | AccountsControllerAccountRemovedEvent | AccountsControllerAccountBalancesUpdatesEvent - | MultichainAssetsControllerStateChangeEvent; - + | MultichainAssetsControllerAccountAssetListUpdatedEvent; /** * Messenger type for the MultichainBalancesController. */ @@ -174,22 +172,77 @@ export class MultichainBalancesController extends BaseController< (balanceUpdate: AccountBalancesUpdatedEventPayload) => this.#handleOnAccountBalancesUpdated(balanceUpdate), ); - // TODO: Maybe add a MultichainAssetsController:accountAssetListUpdated event instead of using the entire state. - // Since MultichainAssetsController already listens for the AccountsController:accountAdded, we can rely in it for that event - // and not listen for it also here, in this controller, since it would be redundant + this.messagingSystem.subscribe( - 'MultichainAssetsController:stateChange', - async (assetsState: MultichainAssetsControllerState) => { - for (const accountId of Object.keys(assetsState.accountsAssets)) { - await this.#updateBalance( + 'MultichainAssetsController:accountAssetListUpdated', + async ({ assets }) => { + const newAccountAssets = Object.entries(assets).map( + ([accountId, { added }]) => ({ accountId, - assetsState.accountsAssets[accountId], - ); - } + assets: [...added], + }), + ); + await this.#handleOnAccountAssetListUpdated(newAccountAssets); }, ); } + /** + * Updates the balances for the given accounts. + * + * @param accounts - The accounts to update the balances for. + */ + async #handleOnAccountAssetListUpdated( + accounts: { + accountId: string; + assets: CaipAssetType[]; + }[], + ): Promise { + const { isUnlocked } = this.messagingSystem.call( + 'KeyringController:getState', + ); + + if (!isUnlocked) { + return; + } + const balancesToUpdate: MultichainBalancesControllerState['balances'] = {}; + + for (const { accountId, assets } of accounts) { + const account = this.#getAccount(accountId); + if (account.metadata.snap) { + const accountBalance = await this.#getBalances( + account.id, + account.metadata.snap.id, + assets, + ); + balancesToUpdate[accountId] = accountBalance; + } + } + + if (Object.keys(balancesToUpdate).length === 0) { + return; + } + + this.update((state: Draft) => { + for (const [accountId, accountBalances] of Object.entries( + balancesToUpdate, + )) { + if ( + !state.balances[accountId] || + Object.keys(state.balances[accountId]).length === 0 + ) { + state.balances[accountId] = accountBalances; + } else { + for (const assetId in accountBalances) { + if (!state.balances[accountId][assetId]) { + state.balances[accountId][assetId] = accountBalances[assetId]; + } + } + } + } + }); + } + /** * Updates the balances of one account. This method doesn't return * anything, but it updates the state of the controller. @@ -262,7 +315,6 @@ export class MultichainBalancesController extends BaseController< */ #listAccounts(): InternalAccount[] { const accounts = this.#listMultichainAccounts(); - return accounts.filter((account) => this.#isNonEvmAccount(account)); } diff --git a/packages/assets-controllers/src/index.ts b/packages/assets-controllers/src/index.ts index d2fd8c89a62..cd4075ec4b7 100644 --- a/packages/assets-controllers/src/index.ts +++ b/packages/assets-controllers/src/index.ts @@ -169,6 +169,7 @@ export type { MultichainAssetsControllerStateChangeEvent, MultichainAssetsControllerActions, MultichainAssetsControllerEvents, + MultichainAssetsControllerAccountAssetListUpdatedEvent, MultichainAssetsControllerMessenger, } from './MultichainAssetsController';