diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 290147ce5c1..0e5e2131fbd 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- `AccountActivityDataSource` is now the highest-priority balance data source and participates in chain-claiming: chains it reports as "up" (from `AccountActivityService:statusChanged`) are claimed first so the polling data sources (`AccountsApiDataSource`/`RpcDataSource`) do not also poll them ([#9517](https://github.com/MetaMask/core/pull/9517)) + - `AssetsController` no longer references `BackendWebSocketService` actions/events; real-time balances and chain status are consumed exclusively from `AccountActivityService` + +### Removed + +- **BREAKING:** Remove `BackendWebsocketDataSource` and its factory/types (`BackendWebsocketDataSource`, `createBackendWebsocketDataSource`, `BackendWebsocketDataSourceOptions`, `BackendWebsocketDataSourceState`). Real-time balance updates and per-chain status are now consumed from `AccountActivityService` via `AccountActivityDataSource`, which manages the WebSocket connection and subscriptions. Consumers no longer need to delegate `BackendWebSocketService` actions/events to the `AssetsController` messenger ([#9517](https://github.com/MetaMask/core/pull/9517)) + +### Fixed + +- The coalesced re-subscribe scheduled by `AssetsController` now re-checks the run guards (UI open, keyring unlocked, and `isEnabled`) when its debounce/jitter timer fires. Previously a re-subscribe scheduled while tracking was allowed could still run after the controller was disabled or tracking was stopped (e.g. UI closed or keyring locked), inappropriately restarting polling subscriptions ([#9517](https://github.com/MetaMask/core/pull/9517)) + ## [11.2.1] ### Fixed @@ -55,11 +68,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `TokensApiClient` (used by `RpcDataSource` / `TokenDetector`) now sets the token-list `occurrenceFloor` query param from the same Token API `GET /v1/suggestedOccurrenceFloors` endpoint (cached 1h), replacing the hardcoded Linea/MegaETH/Tempo special cases. Missing chains or failed fetches fall back to 3 ([#9537](https://github.com/MetaMask/core/pull/9537)) - Bump `@metamask/network-enablement-controller` from `^5.5.0` to `^5.6.0` ([#9520](https://github.com/MetaMask/core/pull/9520)) - Bump `@metamask/phishing-controller` from `^17.2.1` to `^17.3.0` ([#9532](https://github.com/MetaMask/core/pull/9532)) -- Bump `@metamask/transaction-controller` from `^69.0.0` to `^69.1.0` ([#9568](https://github.com/MetaMask/core/pull/9568)) - -### Fixed - -- `TokenDataSource` now also fetches token metadata for assets present in `assetsBalance` that are missing `assetsInfo` (previously only detected assets without metadata were enriched) ([#9547](https://github.com/MetaMask/core/pull/9547)) ## [11.0.0] diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 3afadff79af..ca18a6048a9 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -90,6 +90,12 @@ type AllEvents = MessengerEvents; type RootMessenger = Messenger; +/** Mirrors the private `RESUBSCRIBE_DEBOUNCE_MS` in AssetsController.ts. */ +const RESUBSCRIBE_DEBOUNCE_MS = 250; + +/** Mirrors the private `RESUBSCRIBE_JITTER_MS` in AssetsController.ts. */ +const RESUBSCRIBE_JITTER_MS = 5000; + const MOCK_ACCOUNT_ID = 'mock-account-id-1'; const MOCK_ASSET_ID = 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; @@ -97,6 +103,22 @@ const MOCK_ASSET_ID_LOWERCASE = 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48' as Caip19AssetId; const MOCK_NATIVE_ASSET_ID = 'eip155:1/slip44:60' as Caip19AssetId; +/** + * Activate asset tracking by marking the UI open and the keyring unlocked, + * then flushing the async startup so the controller is in its running state. + * + * @param messenger - The root messenger used to publish lifecycle events. + */ +async function activateTracking(messenger: RootMessenger): Promise { + ( + messenger as unknown as { + publish: (topic: string, payload?: unknown) => void; + } + ).publish('ClientController:stateChange', { isUiOpen: true }); + messenger.publish('KeyringController:unlock'); + await flushPromises(); +} + function createMockInternalAccount( overrides?: Partial, ): InternalAccount { @@ -740,7 +762,7 @@ describe('AssetsController', () => { }); }); - it('graduates an EVM custom asset when BackendWebsocketDataSource reports a balance for it', async () => { + it('graduates an EVM custom asset when AccountActivityDataSource reports a balance for it', async () => { await withController(async ({ controller }) => { await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); @@ -752,7 +774,7 @@ describe('AssetsController', () => { }, }, }, - 'BackendWebsocketDataSource', + 'AccountActivityDataSource', ); expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined(); @@ -1608,14 +1630,99 @@ describe('AssetsController', () => { }); describe('handleActiveChainsUpdate', () => { - it('re-subscribes assets when chains are added', async () => { - await withController(async ({ controller }) => { - const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + it('re-subscribes assets when chains are added, debounced and jittered', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.5); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); - const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); - onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + + // Re-subscribe is debounced, so it does not run synchronously. + expect(subscribeSpy).not.toHaveBeenCalled(); + + // Additions are jittered on top of the base debounce, so nothing runs + // at the base debounce window. + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_MS); + expect(subscribeSpy).not.toHaveBeenCalled(); + + // Once the jitter window elapses it runs exactly once. + jest.advanceTimersByTime(RESUBSCRIBE_JITTER_MS); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } + }); + }); + + it('lets a shorter-delay re-subscribe pre-empt a pending longer-delay one', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + // First update draws a long jitter, second update draws a short one. + const randomSpy = jest + .spyOn(Math, 'random') + .mockReturnValueOnce(0.9) + .mockReturnValueOnce(0.1); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); - expect(subscribeSpy).toHaveBeenCalledTimes(1); + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + // Schedules a long-delay jittered re-subscribe (~4750ms)... + onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + // ...then a second update with a shorter delay (~750ms) pre-empts it. + onActiveChainsUpdated('TestDataSource', [], ['eip155:137']); + + // Runs once at the shorter window. + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_MS + 500); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + + // The pre-empted longer-delay timer does not fire a second time. + jest.advanceTimersByTime(RESUBSCRIBE_JITTER_MS); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } + }); + }); + + it('coalesces a burst of active-chain updates into a single re-subscribe', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.5); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + // Simulate a burst of chain up/down notifications within the window. + onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + onActiveChainsUpdated( + 'TestDataSource', + ['eip155:1', 'eip155:137'], + ['eip155:1'], + ); + onActiveChainsUpdated( + 'TestDataSource', + ['eip155:137'], + ['eip155:1', 'eip155:137'], + ); + + // Advance past the base debounce plus the jitter window; the burst + // collapses into exactly one re-subscribe. + jest.advanceTimersByTime( + RESUBSCRIBE_DEBOUNCE_MS + RESUBSCRIBE_JITTER_MS, + ); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } }); }); @@ -1642,14 +1749,29 @@ describe('AssetsController', () => { }); }); - it('re-subscribes assets when chains are removed', async () => { - await withController(async ({ controller }) => { - const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + it('re-subscribes assets when chains are removed, debounced and jittered', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.5); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); - const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); - onActiveChainsUpdated('TestDataSource', [], ['eip155:1']); + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + onActiveChainsUpdated('TestDataSource', [], ['eip155:1']); + + // Removals are jittered on top of the base debounce, so nothing runs + // at the base debounce window. + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_MS); + expect(subscribeSpy).not.toHaveBeenCalled(); - expect(subscribeSpy).toHaveBeenCalledTimes(1); + // Once the jitter window elapses it runs exactly once. + jest.advanceTimersByTime(RESUBSCRIBE_JITTER_MS); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } }); }); @@ -1687,6 +1809,68 @@ describe('AssetsController', () => { }, ); }); + + it('does not run a scheduled re-subscribe when the controller is disabled before the timer fires', async () => { + let enabled = true; + + await withController( + { + controllerOptions: { isEnabled: (): boolean => enabled }, + }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.5); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + + // Scheduled while enabled... + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + + // ...but the controller is disabled before the timer fires. + enabled = false; + + jest.advanceTimersByTime( + RESUBSCRIBE_DEBOUNCE_MS + RESUBSCRIBE_JITTER_MS, + ); + + expect(subscribeSpy).not.toHaveBeenCalled(); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } + }, + ); + }); + + it('does not run a scheduled re-subscribe when the keyring locks before the timer fires', async () => { + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.5); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + + // The keyring locks (e.g. the app is backgrounded). A data source can + // still report chains going "down" as the WebSocket disconnects, + // which schedules a coalesced re-subscribe even though tracking has + // been torn down. + messenger.publish('KeyringController:lock'); + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + onActiveChainsUpdated('TestDataSource', [], ['eip155:1']); + + jest.advanceTimersByTime( + RESUBSCRIBE_DEBOUNCE_MS + RESUBSCRIBE_JITTER_MS, + ); + + expect(subscribeSpy).not.toHaveBeenCalled(); + } finally { + randomSpy.mockRestore(); + jest.useRealTimers(); + } + }); + }); }); describe('handleAssetsUpdate - state updates', () => { @@ -2235,8 +2419,19 @@ describe('AssetsController', () => { }; await withController( - { state: initialState }, + { state: initialState, clientControllerState: { isUiOpen: true } }, async ({ controller, messenger }) => { + // UI must be open and keyring unlocked so AccountActivityDataSource is + // subscribed and can route balanceUpdated events to the account. + ( + messenger as unknown as { + publish: (topic: string, payload?: unknown) => void; + } + ).publish('ClientController:stateChange', { isUiOpen: true }); + messenger.publish('KeyringController:unlock'); + + await flushPromises(); + messenger.publish('AccountActivityService:balanceUpdated', { address: '0x1234567890123456789012345678901234567890', chain: 'eip155:42161', diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 3fb7d56a85e..8b636000871 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -16,9 +16,7 @@ import type { TraceCallback } from '@metamask/controller-utils'; import type { ApiPlatformClient, AccountActivityServiceBalanceUpdatedEvent, - BackendWebSocketServiceActions, - BackendWebSocketServiceEvents, - BalanceUpdate, + AccountActivityServiceStatusChangedEvent, SupportedCurrency, } from '@metamask/core-backend'; import type { @@ -79,9 +77,9 @@ import type { DataSourceState, SubscriptionRequest, } from './data-sources/AbstractDataSource.js'; +import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource.js'; import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource.js'; import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource.js'; -import { BackendWebsocketDataSource } from './data-sources/BackendWebsocketDataSource.js'; import { shouldSkipNativeForCaipChainId } from './data-sources/evm-rpc-services/utils/assets.js'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource.js'; import { @@ -153,7 +151,6 @@ import type { BridgeExchangeRatesFormat, TransactionPayLegacyFormat, } from './utils/index.js'; -import { processAccountActivityBalanceUpdates } from './utils/processAccountActivityBalanceUpdates.js'; const NATIVE_ASSETS_QUERY_KEY = ['nativeAssets']; @@ -206,6 +203,25 @@ const MESSENGER_EXPOSED_METHODS = [ /** Default polling interval hint for data sources (30 seconds) */ const DEFAULT_POLLING_INTERVAL_MS = 30_000; +/** + * Trailing debounce window (ms) for coalescing event-driven re-subscribes. + * A burst of `onActiveChainsUpdated` callbacks (e.g. many `statusChanged` + * up/down notifications, or several data sources reporting supported chains at + * init) collapses into a single re-subscribe pass so polling data sources are + * not torn down and re-created (with a fresh immediate fetch) once per event. + */ +const RESUBSCRIBE_DEBOUNCE_MS = 250; + +/** + * Maximum random jitter (ms) added to the re-subscribe delay whenever a data + * source's active chains change (e.g. WebSocket `statusChanged` up/down). + * Staggers the resulting WS-subscribe fan-out across clients that all receive + * the same backend broadcast, avoiding a thundering herd. Applied to both chain + * additions and removals; the earliest scheduled run still wins so a + * shorter-delay request pre-empts a pending longer-delay one. + */ +const RESUBSCRIBE_JITTER_MS = 5000; + // ============================================================================ // TRACE NAMES — used in Sentry spans (search these strings in Discover) // ============================================================================ @@ -328,8 +344,6 @@ type AllowedActions = | SnapControllerGetRunnableSnapsAction | SnapControllerHandleRequestAction | GetPermissions - // BackendWebsocketDataSource - | BackendWebSocketServiceActions // PhishingController | PhishingControllerBulkScanTokensAction // AccountsApiDataSource (Accounts API v6 balances feature flag) @@ -359,10 +373,9 @@ type AllowedEvents = | AccountsControllerAccountBalancesUpdatedEvent | PermissionControllerStateChange | SnapControllerSnapInstalledEvent - // BackendWebsocketDataSource - | BackendWebSocketServiceEvents - // AccountActivityService (real-time balance updates for unified assets) + // AccountActivityService (real-time balance updates + chain status for unified assets) | AccountActivityServiceBalanceUpdatedEvent + | AccountActivityServiceStatusChangedEvent // AccountsApiDataSource subscribes to react to Snaps → AssetsController // migration flag changes (which gate the chains it surfaces as active) | RemoteFeatureFlagControllerStateChangeEvent; @@ -757,6 +770,20 @@ export class AssetsController extends BaseController< */ readonly #activeSubscriptions: Map = new Map(); + /** + * Pending trailing-debounce timer for coalesced re-subscribes. Bursts of + * event-driven `#scheduleSubscribeAssets()` calls collapse into a single + * `#subscribeAssets()` run. See {@link RESUBSCRIBE_DEBOUNCE_MS}. + */ + #resubscribeTimer: ReturnType | null = null; + + /** + * Scheduled fire time (epoch ms) of {@link #resubscribeTimer}. Lets an urgent + * (non-jittered) re-subscribe pre-empt a pending jittered one so a chain + * removal is not delayed behind a chain addition's jitter window. + */ + #resubscribeRunAt = 0; + /** Currently enabled chains from NetworkEnablementController */ #enabledChains: Set = new Set(); @@ -790,7 +817,7 @@ export class AssetsController extends BaseController< return []; } - readonly #backendWebsocketDataSource: BackendWebsocketDataSource; + readonly #accountActivityDataSource: AccountActivityDataSource; readonly #accountsApiDataSource: AccountsApiDataSource; @@ -802,19 +829,25 @@ export class AssetsController extends BaseController< /** * All balance data sources in priority order for chain-claiming and cleanup. - * Note: StakedBalanceDataSource is excluded because it provides supplementary - * data and should not participate in chain-claiming. + * AccountActivityDataSource is highest priority: chains it reports as "up" + * (real-time WebSocket data via AccountActivityService) are reserved first so + * the polling sources (AccountsApi/RPC) do not also poll them. It only + * reserves chains here — its actual subscription (used to route + * `balanceUpdated` events) is created separately in `#subscribeAssetsBalance` + * with all selected accounts and enabled chains. Note: StakedBalanceDataSource + * is excluded because it provides supplementary data and should not + * participate in chain-claiming. * * @returns The four balance data source instances in priority order. */ get #allBalanceDataSources(): [ - BackendWebsocketDataSource, + AccountActivityDataSource, AccountsApiDataSource, SnapDataSource, RpcDataSource, ] { return [ - this.#backendWebsocketDataSource, + this.#accountActivityDataSource, this.#accountsApiDataSource, this.#snapDataSource, this.#rpcDataSource, @@ -905,12 +938,17 @@ export class AssetsController extends BaseController< } }; - this.#backendWebsocketDataSource = new BackendWebsocketDataSource({ + this.#accountActivityDataSource = new AccountActivityDataSource({ messenger: this.messenger, - queryApiClient, onActiveChainsUpdated: this.#onActiveChainsUpdated, getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => this.#getAssetType(assetId), + onAssetsUpdate: (response, request): Promise => + this.handleAssetsUpdate( + response, + this.#accountActivityDataSource.getName(), + request, + ), }); this.#accountsApiDataSource = new AccountsApiDataSource({ messenger: this.messenger, @@ -1188,16 +1226,6 @@ export class AssetsController extends BaseController< this.#onTransactionConfirmed(transactionMeta); }, ); - - // Real-time post-tx balances from AccountActivityService (same WS payload as - // TokenBalancesController; BackendWebsocketDataSource may not receive the - // callback when AccountActivityService owns the server subscription). - this.messenger.subscribe( - 'AccountActivityService:balanceUpdated', - (event) => { - this.#onAccountActivityBalanceUpdated(event); - }, - ); } #onUnapprovedTransactionAdded(transactionMeta: TransactionMeta): void { @@ -1256,49 +1284,6 @@ export class AssetsController extends BaseController< }); } - #onAccountActivityBalanceUpdated({ - address, - chain, - updates, - }: { - address: string; - chain: string; - updates: BalanceUpdate[]; - }): void { - const account = this.#getSelectedAccounts().find((a) => - a.address.startsWith('0x') - ? a.address.toLowerCase() === address.toLowerCase() - : a.address === address, - ); - - if (!account) { - return; - } - - const chainId = chain as ChainId; - const response = processAccountActivityBalanceUpdates( - updates, - account.id, - (assetId) => this.#getAssetType(assetId), - ); - - if (!response.assetsBalance) { - return; - } - - const request: DataRequest = { - accountsWithSupportedChains: [{ account, supportedChains: [chainId] }], - chainIds: [chainId], - dataTypes: ['balance', 'metadata'], - }; - - this.handleAssetsUpdate(response, 'AccountActivityService', request).catch( - (error) => { - log('Failed to apply AccountActivityService balance update', { error }); - }, - ); - } - /** * Start or stop asset tracking based on client (UI) open state and keyring * unlock state. Only runs when both UI is open and keyring is unlocked. @@ -1465,12 +1450,18 @@ export class AssetsController extends BaseController< const addedChains = activeChains.filter((ch) => !previousSet.has(ch)); const removedChains = previous.filter((ch) => !activeChains.includes(ch)); + // Refresh subscriptions to use updated data source availability. + // No one-time fetch needed here — #handleEnabledNetworksChanged + // handles fetches when the user enables a network, and + // #subscribeAssets re-subscribes with the correct chain assignment. + // Coalesce bursts of chain updates into a single re-subscribe so polling + // data sources are not torn down/re-created (with a redundant immediate + // fetch) once per event. + // Jitter is always applied to stagger the resulting WS-subscribe fan-out + // across clients receiving the same backend broadcast, avoiding a + // thundering herd for both chain additions and removals. if (addedChains.length > 0 || removedChains.length > 0) { - // Refresh subscriptions to use updated data source availability. - // No one-time fetch needed here — #handleEnabledNetworksChanged - // handles fetches when the user enables a network, and - // #subscribeAssets re-subscribes with the correct chain assignment. - this.#subscribeAssets(); + this.#scheduleSubscribeAssets(); } } @@ -2965,6 +2956,13 @@ export class AssetsController extends BaseController< this.#stateSizeReported = false; this.#lastKnownAccountIds = new Set(); + // Cancel any pending coalesced re-subscribe so it cannot fire after stop. + if (this.#resubscribeTimer) { + clearTimeout(this.#resubscribeTimer); + this.#resubscribeTimer = null; + } + this.#resubscribeRunAt = 0; + // Stop price subscription first (uses direct messenger call) this.unsubscribeAssetsPrice(); @@ -3010,6 +3008,67 @@ export class AssetsController extends BaseController< this.#subscribeAssets(); } + /** + * Schedule a coalesced re-subscribe. Multiple calls within + * {@link RESUBSCRIBE_DEBOUNCE_MS} collapse into a single `#subscribeAssets()` + * run. Used for event-driven bursts (e.g. `onActiveChainsUpdated` from many + * `statusChanged` notifications) so polling data sources are not repeatedly + * torn down and re-created (each re-create firing an immediate, redundant + * fetch). Explicit flows (startup, account changes, enabled-network changes) + * keep calling `#subscribeAssets()` directly so their re-subscribe is not + * delayed. + * + * A random delay up to {@link RESUBSCRIBE_JITTER_MS} is added on top of the + * base debounce to stagger the WS-subscribe fan-out across clients receiving + * the same backend broadcast, avoiding a thundering herd. The earliest + * scheduled run always wins, so a shorter-delay request pre-empts a pending + * longer-delay one while a later request never pushes a sooner pending run out. + */ + #scheduleSubscribeAssets(): void { + const delay = + RESUBSCRIBE_DEBOUNCE_MS + + Math.floor(Math.random() * RESUBSCRIBE_JITTER_MS); + const runAt = Date.now() + delay; + + if (this.#resubscribeTimer) { + // Keep the earliest scheduled run so an urgent (shorter-delay) request can + // pre-empt a pending jittered one, but a jittered request never pushes a + // sooner pending run later. + if (runAt >= this.#resubscribeRunAt) { + return; + } + clearTimeout(this.#resubscribeTimer); + this.#resubscribeTimer = null; + } + + this.#resubscribeRunAt = runAt; + this.#resubscribeTimer = setTimeout(() => { + this.#resubscribeTimer = null; + this.#resubscribeRunAt = 0; + // Re-check the run guards at fire time. A re-subscribe can be scheduled + // while tracking is allowed but fire after it should have stopped: + // #handleActiveChainsUpdate only gates on #isEnabled(), so a data source + // reporting chains "down" (e.g. on WebSocket disconnect) after the UI + // closed or the keyring locked can still schedule this timer; and + // #isEnabled() is an external getter that can flip without routing + // through #stop() (which is what cancels the timer). Running + // #subscribeAssets() regardless would inappropriately restart polling. + if (!this.#uiOpen || !this.#keyringUnlocked || !this.#isEnabled()) { + log('Skipping coalesced re-subscribe: tracking no longer active', { + uiOpen: this.#uiOpen, + keyringUnlocked: this.#keyringUnlocked, + isEnabled: this.#isEnabled(), + }); + return; + } + try { + this.#subscribeAssets(); + } catch (error) { + log('Failed to run coalesced re-subscribe', error); + } + }, delay); + } + /** * Subscribe to asset updates for all selected accounts. */ @@ -3053,8 +3112,10 @@ export class AssetsController extends BaseController< new Set(chainIds), ); const remainingChains = new Set(chainToAccounts.keys()); - // When basic functionality is on, use all balance data sources; when off, RPC only. - const balanceDataSources = this.#isBasicFunctionality() + // When basic functionality is on, use all balance data sources; when off, + // RPC only. + const isBasicFunctionality = this.#isBasicFunctionality(); + const balanceDataSources = isBasicFunctionality ? this.#allBalanceDataSources : [this.#rpcDataSource]; @@ -3110,7 +3171,7 @@ export class AssetsController extends BaseController< /** * Guarantee that customAssets are **always** polled by RPC, even when - * AccountsApi or the websocket data source has claimed the chain in the + * AccountsApi or another data source has claimed the chain in the * regular handoff. RPC is the sole balance fetcher for user-imported * tokens (see `pickRpcCustomAssetsSupplement` for the full rationale), * so we run a dedicated subscription in `customAssetsOnly` mode under a @@ -3536,14 +3597,11 @@ export class AssetsController extends BaseController< } /** - * Refresh data-source `activeChains` after an EVM network switch so API/WS/Rpc + * Refresh data-source `activeChains` after an EVM network switch so API/Rpc * chain claiming is not stuck on an empty or stale init-time list. */ async #refreshActiveChainsOnNetworkSwitch(): Promise { - await Promise.all([ - this.#accountsApiDataSource.refreshActiveChains(), - this.#backendWebsocketDataSource.refreshActiveChains(), - ]); + await this.#accountsApiDataSource.refreshActiveChains(); this.#rpcDataSource.refreshActiveChainsFromNetworkState(); } @@ -3701,13 +3759,13 @@ export class AssetsController extends BaseController< ), }; - // Graduate custom assets only when AccountsAPI / Websocket reports them. - // RPC already fetches custom assets on purpose, and Snap handles non-EVM - // chains the rule does not apply to, so skip the middleware for those. + // Graduate custom assets only when AccountsAPI / AccountActivity reports + // them. RPC already fetches custom assets on purpose, and Snap handles + // non-EVM chains the rule does not apply to, so skip the middleware for + // those. const shouldGraduateCustomAssets = sourceId === 'AccountsApiDataSource' || - sourceId === 'BackendWebsocketDataSource' || - sourceId === 'AccountActivityService'; + sourceId === 'AccountActivityDataSource'; const enrichmentSources: AssetsDataSource[] = [ ...(shouldGraduateCustomAssets @@ -3758,7 +3816,7 @@ export class AssetsController extends BaseController< }); // Destroy instantiated data sources - this.#backendWebsocketDataSource?.destroy?.(); + this.#accountActivityDataSource?.destroy?.(); this.#accountsApiDataSource?.destroy?.(); this.#snapDataSource?.destroy?.(); this.#rpcDataSource?.destroy?.(); diff --git a/packages/assets-controller/src/README.md b/packages/assets-controller/src/README.md index baebf8606fa..849f915ffe0 100644 --- a/packages/assets-controller/src/README.md +++ b/packages/assets-controller/src/README.md @@ -31,7 +31,7 @@ The `AssetsController` is a unified asset management system that provides real-t │ On-demand data fetch Process incoming updates │ │ (forceUpdate: true) (enrichment only) │ │ │ -│ Subscription Sources: BackendWebsocket, AccountsApi, Snap, RPC │ +│ Subscription Sources: AccountActivity, AccountsApi, Snap, RPC │ │ Push updates via: AssetsController:assetsUpdate action │ └──────────────────────────────────────────────────────────────────────────────┘ ``` @@ -90,7 +90,7 @@ registerActionHandlers() #### 1.5 Balance data source priority -Built-in balance data sources are fixed and processed in priority order: BackendWebsocketDataSource, AccountsApiDataSource, SnapDataSource, RpcDataSource. Earlier sources get first pick for chain assignment; later sources act as fallbacks. Data sources report active chains via the `onActiveChainsUpdated` callback passed at construction. +Built-in balance data sources are fixed and processed in priority order: AccountActivityDataSource, AccountsApiDataSource, SnapDataSource, RpcDataSource. Earlier sources get first pick for chain assignment; later sources act as fallbacks. Data sources report active chains via the `onActiveChainsUpdated` callback passed at construction. #### 1.6 Middleware Chains @@ -506,7 +506,7 @@ messenger.call( // Data source reports its active chains messenger.call( 'AssetsController:activeChainsUpdate', - 'BackendWebsocketDataSource', + 'AccountActivityDataSource', ['eip155:1', 'eip155:137', 'eip155:42161'], ); ``` @@ -538,7 +538,7 @@ await messenger.call( }, }, }, - 'BackendWebsocketDataSource', + 'AccountActivityDataSource', ); ``` @@ -795,12 +795,12 @@ const multiChainAssets = await messenger.call( ### Primary Data Sources (Balance Providers) -| Order | Data Source | Update Mechanism | Chains | -| ----- | -------------------------- | ------------------------ | --------------------- | -| 1 | BackendWebsocketDataSource | Real-time WebSocket push | API-supported EVM | -| 2 | AccountsApiDataSource | HTTP polling | API-supported chains | -| 3 | SnapDataSource | Snap keyring events | Solana, Bitcoin, Tron | -| 4 | RpcDataSource | Direct RPC polling | Any EVM chain | +| Order | Data Source | Update Mechanism | Chains | +| ----- | ------------------------- | ------------------------ | --------------------- | +| 1 | AccountActivityDataSource | Real-time WebSocket push | API-supported EVM | +| 2 | AccountsApiDataSource | HTTP polling | API-supported chains | +| 3 | SnapDataSource | Snap keyring events | Solana, Bitcoin, Tron | +| 4 | RpcDataSource | Direct RPC polling | Any EVM chain | **Fallback behavior**: Data sources are processed in registration order. Earlier sources get first pick for chain assignment; later sources handle remaining chains. @@ -831,11 +831,11 @@ flowchart TB subgraph AssetsControllerInit["AssetsController (with queryApiClient)"] AC[AssetsController] - AC --> WS & API & SNAP & RPC & TOK & PRICE & DET + AC --> AA & API & SNAP & RPC & TOK & PRICE & DET end subgraph DataSources["Data Source Instances - Order 1-4"] - WS[BackendWebsocketDS - Order 1] + AA[AccountActivityDS - Order 1] API[AccountsApiDS - Order 2] SNAP[SnapDataSource - Order 3] RPC[RpcDataSource - Order 4] @@ -849,7 +849,7 @@ flowchart TB ATC[AccountTreeController] NEC[NetworkEnablementController] KC[KeyringController] - BWSS[BackendWebSocketService] + AAS[AccountActivityService] BAC[BackendApiClient / ApiPlatformClient] SC[SnapController] end @@ -858,7 +858,7 @@ flowchart TB ACI --> AC RPC -.-> NC - WS -.-> BWSS + AA -.-> AAS API -.-> BAC SNAP -.-> SC TOK -.-> BAC diff --git a/packages/assets-controller/src/__fixtures__/MockAssetControllerMessenger.ts b/packages/assets-controller/src/__fixtures__/MockAssetControllerMessenger.ts index 4866eb25d5f..fa9c9e3f84e 100644 --- a/packages/assets-controller/src/__fixtures__/MockAssetControllerMessenger.ts +++ b/packages/assets-controller/src/__fixtures__/MockAssetControllerMessenger.ts @@ -61,20 +61,6 @@ export function createMockAssetControllerMessenger(): { 'SnapController:getRunnableSnaps', 'SnapController:handleRequest', 'PermissionController:getPermissions', - // BackendWebsocketDataSource - 'BackendWebSocketService:connect', - 'BackendWebSocketService:disconnect', - 'BackendWebSocketService:forceReconnection', - 'BackendWebSocketService:sendMessage', - 'BackendWebSocketService:sendRequest', - 'BackendWebSocketService:getConnectionInfo', - 'BackendWebSocketService:getSubscriptionsByChannel', - 'BackendWebSocketService:channelHasSubscription', - 'BackendWebSocketService:findSubscriptionsByChannelPrefix', - 'BackendWebSocketService:addChannelCallback', - 'BackendWebSocketService:removeChannelCallback', - 'BackendWebSocketService:getChannelCallbacks', - 'BackendWebSocketService:subscribe', ], events: [ // AssetsController @@ -92,10 +78,9 @@ export function createMockAssetControllerMessenger(): { // SnapDataSource 'AccountsController:accountBalancesUpdated', 'PermissionController:stateChange', - // BackendWebsocketDataSource - 'BackendWebSocketService:connectionStateChanged', - // AccountActivityService + // AccountActivityService (real-time balances + chain status) 'AccountActivityService:balanceUpdated', + 'AccountActivityService:statusChanged', ], }); diff --git a/packages/assets-controller/src/data-sources/AccountActivityDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountActivityDataSource.test.ts new file mode 100644 index 00000000000..293814e6c34 --- /dev/null +++ b/packages/assets-controller/src/data-sources/AccountActivityDataSource.test.ts @@ -0,0 +1,747 @@ +// `getActiveChainsSync` is part of the data source's public API; asserting on +// it is intentional and not a filesystem-style sync call. +/* eslint-disable n/no-sync */ +import type { BalanceUpdate } from '@metamask/core-backend'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { + MessengerActions, + MessengerEvents, + MockAnyNamespace, +} from '@metamask/messenger'; +import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; + +import type { AssetsControllerMessenger } from '../AssetsController.js'; +import type { ChainId, Caip19AssetId } from '../types.js'; +import { + AccountActivityDataSource, + createAccountActivityDataSource, +} from './AccountActivityDataSource.js'; + +type RootMessenger = Messenger< + MockAnyNamespace, + MessengerActions, + MessengerEvents +>; + +const CHAIN_MAINNET = 'eip155:1' as ChainId; +const CHAIN_POLYGON = 'eip155:137' as ChainId; +const CHAIN_SOLANA = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; + +const EVM_ADDRESS = '0x1234567890123456789012345678901234567890'; +const SOLANA_ADDRESS = 'DjVE6JNiYqPL2QXyCUUh8rNjHrbz9hXHNYt99MQ59qw1'; + +const ETH_ASSET = 'eip155:1/slip44:60' as Caip19AssetId; + +/** + * Build an InternalAccount for tests. + * + * @param overrides - Partial fields to override on the account. + * @returns A mock InternalAccount. + */ +function createMockAccount( + overrides?: Partial, +): InternalAccount { + return { + id: 'mock-account-id', + address: EVM_ADDRESS, + options: {}, + methods: [], + type: 'eip155:eoa', + scopes: ['eip155:0'], + metadata: { + name: 'Test Account', + keyring: { type: 'HD Key Tree' }, + importTime: Date.now(), + lastSelected: Date.now(), + }, + ...overrides, + } as InternalAccount; +} + +/** + * Build a BalanceUpdate for tests. + * + * @param overrides - Partial fields to override on the update. + * @param overrides.asset - Partial asset fields to override. + * @param overrides.postBalance - Partial post-balance fields to override. + * @param overrides.transfers - Transfers to set on the update. + * @returns A mock BalanceUpdate. + */ +function createBalanceUpdate(overrides?: { + asset?: Partial; + postBalance?: Partial; + transfers?: BalanceUpdate['transfers']; +}): BalanceUpdate { + return { + asset: { + fungible: true, + type: ETH_ASSET, + unit: 'ETH', + decimals: 18, + ...overrides?.asset, + }, + postBalance: { + amount: '1000000000000000000', + ...overrides?.postBalance, + }, + transfers: overrides?.transfers ?? [], + } as BalanceUpdate; +} + +type SetupOptions = { + groupAccounts?: InternalAccount[]; + selectedAccount?: InternalAccount | null; + getAssetType?: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl'; + onAssetsUpdate?: jest.Mock; + onActiveChainsUpdated?: jest.Mock; + state?: { activeChains?: ChainId[] }; +}; + +type SetupResult = { + dataSource: AccountActivityDataSource; + rootMessenger: RootMessenger; + onAssetsUpdate: jest.Mock; + onActiveChainsUpdated: jest.Mock; + getAssetType: jest.Mock; + triggerBalanceUpdated: (payload: { + address: string; + chain: string; + updates: BalanceUpdate[]; + }) => void; + triggerStatusChanged: (payload: { + chainIds: string[]; + status: 'up' | 'down'; + timestamp?: number; + }) => void; + cleanup: () => void; +}; + +/** + * Create an AccountActivityDataSource wired to a messenger for testing. + * + * @param options - Setup overrides. + * @returns The data source and its test harness. + */ +function setup(options: SetupOptions = {}): SetupResult { + const { + groupAccounts = [createMockAccount()], + selectedAccount = null, + onAssetsUpdate = jest.fn().mockResolvedValue(undefined), + onActiveChainsUpdated = jest.fn(), + state, + } = options; + + const rootMessenger: RootMessenger = new Messenger({ + namespace: MOCK_ANY_NAMESPACE, + }); + + const assetsControllerMessenger: AssetsControllerMessenger = new Messenger({ + namespace: 'AssetsController', + parent: rootMessenger, + }); + + rootMessenger.delegate({ + messenger: assetsControllerMessenger, + actions: [ + 'AccountTreeController:getAccountsFromSelectedAccountGroup', + 'AccountsController:getSelectedAccount', + ], + events: [ + 'AccountActivityService:balanceUpdated', + 'AccountActivityService:statusChanged', + ], + }); + + rootMessenger.registerActionHandler( + 'AccountTreeController:getAccountsFromSelectedAccountGroup', + () => groupAccounts, + ); + rootMessenger.registerActionHandler( + 'AccountsController:getSelectedAccount', + () => selectedAccount as InternalAccount, + ); + + const getAssetType = jest + .fn() + .mockImplementation(options.getAssetType ?? ((): 'native' => 'native')); + + const dataSource = new AccountActivityDataSource({ + messenger: assetsControllerMessenger, + onActiveChainsUpdated, + getAssetType, + onAssetsUpdate, + state, + }); + + const triggerBalanceUpdated = (payload: { + address: string; + chain: string; + updates: BalanceUpdate[]; + }): void => { + rootMessenger.publish('AccountActivityService:balanceUpdated', payload); + }; + + const triggerStatusChanged = (payload: { + chainIds: string[]; + status: 'up' | 'down'; + timestamp?: number; + }): void => { + rootMessenger.publish('AccountActivityService:statusChanged', payload); + }; + + const cleanup = (): void => { + dataSource.destroy(); + rootMessenger.clearSubscriptions(); + }; + + return { + dataSource, + rootMessenger, + onAssetsUpdate, + onActiveChainsUpdated, + getAssetType, + triggerBalanceUpdated, + triggerStatusChanged, + cleanup, + }; +} + +describe('AccountActivityDataSource', () => { + describe('constructor', () => { + it('uses the default (empty) active chains state', () => { + const { dataSource, cleanup } = setup(); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([]); + + cleanup(); + }); + + it('seeds active chains from provided state', () => { + const { dataSource, cleanup } = setup({ + state: { activeChains: [CHAIN_MAINNET] }, + }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([CHAIN_MAINNET]); + + cleanup(); + }); + + it('subscribes to balanceUpdated and statusChanged events', () => { + const { dataSource, onAssetsUpdate, triggerStatusChanged, cleanup } = + setup(); + + triggerStatusChanged({ chainIds: [CHAIN_MAINNET], status: 'up' }); + + // statusChanged handler ran and claimed the chain + expect(dataSource.getActiveChainsSync()).toStrictEqual([CHAIN_MAINNET]); + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + }); + + describe('getName', () => { + it('returns the data source name', () => { + const { dataSource, cleanup } = setup(); + + expect(dataSource.getName()).toBe('AccountActivityDataSource'); + + cleanup(); + }); + }); + + describe('subscribe', () => { + it('resolves without doing anything (no-op)', async () => { + const { dataSource, onAssetsUpdate, onActiveChainsUpdated, cleanup } = + setup(); + + expect(await dataSource.subscribe()).toBeUndefined(); + expect(onAssetsUpdate).not.toHaveBeenCalled(); + expect(onActiveChainsUpdated).not.toHaveBeenCalled(); + + cleanup(); + }); + }); + + describe('balanceUpdated event', () => { + it('reports decoded balances through onAssetsUpdate', async () => { + const account = createMockAccount(); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [account], + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + const [response, request] = onAssetsUpdate.mock.calls[0]; + // `processAccountActivityBalanceUpdates` builds null-prototype objects + // (`Object.create(null)`), so `toStrictEqual` (which checks prototypes) + // does not match a plain object literal here. + // eslint-disable-next-line jest/prefer-strict-equal + expect(response).toEqual({ + updateMode: 'merge', + assetsBalance: { + [account.id]: { + [ETH_ASSET]: { amount: '1' }, + }, + }, + assetsInfo: { + [ETH_ASSET]: { + type: 'native', + symbol: 'ETH', + name: 'ETH', + decimals: 18, + }, + }, + }); + expect(request).toStrictEqual({ + accountsWithSupportedChains: [ + { account, supportedChains: [CHAIN_MAINNET] }, + ], + chainIds: [CHAIN_MAINNET], + dataTypes: ['balance', 'metadata'], + }); + + cleanup(); + }); + + it('converts a hex postBalance to a human-readable amount', async () => { + const account = createMockAccount(); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [account], + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [ + createBalanceUpdate({ postBalance: { amount: '0x10aa6d94e80' } }), + ], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + const [response] = onAssetsUpdate.mock.calls[0]; + expect(response.assetsBalance[account.id][ETH_ASSET]).toStrictEqual({ + amount: '0.00000114526056', + }); + + cleanup(); + }); + + it('resolves the asset type via the injected getAssetType', async () => { + const { getAssetType, triggerBalanceUpdated, cleanup } = setup({ + getAssetType: () => 'erc20', + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(getAssetType).toHaveBeenCalledWith(ETH_ASSET); + + cleanup(); + }); + + it.each([ + ['address is empty', { address: '' }], + ['chain is empty', { chain: '' }], + ['updates is empty', { updates: [] }], + ])('ignores the event when %s', async (_label, override) => { + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup(); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + ...override, + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('ignores the event when no wallet account matches the address', async () => { + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [ + createMockAccount({ + address: '0x9999999999999999999999999999999999999999', + }), + ], + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('matches EVM addresses case-insensitively', async () => { + const account = createMockAccount({ address: EVM_ADDRESS.toLowerCase() }); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [account], + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS.toUpperCase().replace('0X', '0x'), + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + + cleanup(); + }); + + it('matches non-EVM addresses exactly', async () => { + const account = createMockAccount({ + address: SOLANA_ADDRESS, + type: 'solana:data-account', + }); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [account], + }); + + triggerBalanceUpdated({ + address: SOLANA_ADDRESS, + chain: CHAIN_SOLANA, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + + cleanup(); + }); + + it('does not match a non-EVM address that differs only by case', async () => { + const account = createMockAccount({ + address: SOLANA_ADDRESS, + type: 'solana:data-account', + }); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [account], + }); + + triggerBalanceUpdated({ + address: SOLANA_ADDRESS.toLowerCase(), + chain: CHAIN_SOLANA, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('falls back to the selected account when the group is empty', async () => { + const selected = createMockAccount({ id: 'selected-id' }); + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [], + selectedAccount: selected, + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + const [response] = onAssetsUpdate.mock.calls[0]; + expect(response.assetsBalance).toHaveProperty('selected-id'); + + cleanup(); + }); + + it('ignores the event when the group is empty and there is no selected account', async () => { + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + groupAccounts: [], + selectedAccount: null, + }); + + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('does not call onAssetsUpdate when no balances could be decoded', async () => { + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup(); + + // decimals undefined => the update is skipped and assetsBalance is absent + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [ + createBalanceUpdate({ + asset: { decimals: undefined as unknown as number }, + }), + ], + }); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('swallows rejections from onAssetsUpdate', async () => { + const onAssetsUpdate = jest + .fn() + .mockRejectedValue(new Error('report failed')); + const { triggerBalanceUpdated, cleanup } = setup({ onAssetsUpdate }); + + expect(() => + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }), + ).not.toThrow(); + + await Promise.resolve(); + await Promise.resolve(); + + expect(onAssetsUpdate).toHaveBeenCalledTimes(1); + + cleanup(); + }); + + it('swallows synchronous errors thrown while handling the event', async () => { + const { onAssetsUpdate, triggerBalanceUpdated, cleanup } = setup({ + getAssetType: () => { + throw new Error('boom'); + }, + }); + + expect(() => + triggerBalanceUpdated({ + address: EVM_ADDRESS, + chain: CHAIN_MAINNET, + updates: [createBalanceUpdate()], + }), + ).not.toThrow(); + + await Promise.resolve(); + + expect(onAssetsUpdate).not.toHaveBeenCalled(); + + cleanup(); + }); + }); + + describe('statusChanged event', () => { + it('claims chains reported as up', () => { + const { + dataSource, + onActiveChainsUpdated, + triggerStatusChanged, + cleanup, + } = setup(); + + triggerStatusChanged({ + chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], + status: 'up', + }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([ + CHAIN_MAINNET, + CHAIN_POLYGON, + ]); + expect(onActiveChainsUpdated).toHaveBeenCalledWith( + 'AccountActivityDataSource', + [CHAIN_MAINNET, CHAIN_POLYGON], + [], + ); + + cleanup(); + }); + + it('releases chains reported as down', () => { + const { + dataSource, + onActiveChainsUpdated, + triggerStatusChanged, + cleanup, + } = setup({ state: { activeChains: [CHAIN_MAINNET, CHAIN_POLYGON] } }); + + triggerStatusChanged({ chainIds: [CHAIN_MAINNET], status: 'down' }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([CHAIN_POLYGON]); + expect(onActiveChainsUpdated).toHaveBeenCalledWith( + 'AccountActivityDataSource', + [CHAIN_POLYGON], + [CHAIN_MAINNET, CHAIN_POLYGON], + ); + + cleanup(); + }); + + it('acts on non-EVM namespaces such as solana', () => { + const { dataSource, triggerStatusChanged, cleanup } = setup(); + + triggerStatusChanged({ chainIds: [CHAIN_SOLANA], status: 'up' }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([CHAIN_SOLANA]); + + cleanup(); + }); + + it('filters out identifiers that are not valid CAIP-2 chain IDs', () => { + const { + dataSource, + onActiveChainsUpdated, + triggerStatusChanged, + cleanup, + } = setup(); + + triggerStatusChanged({ + chainIds: ['not-a-chain', CHAIN_MAINNET], + status: 'up', + }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([CHAIN_MAINNET]); + expect(onActiveChainsUpdated).toHaveBeenCalledWith( + 'AccountActivityDataSource', + [CHAIN_MAINNET], + [], + ); + + cleanup(); + }); + + it('does nothing when there are no valid chains', () => { + const { + dataSource, + onActiveChainsUpdated, + triggerStatusChanged, + cleanup, + } = setup(); + + triggerStatusChanged({ chainIds: ['not-a-chain'], status: 'up' }); + + expect(dataSource.getActiveChainsSync()).toStrictEqual([]); + expect(onActiveChainsUpdated).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('does not notify listeners when the active chains do not change', () => { + const { onActiveChainsUpdated, triggerStatusChanged, cleanup } = setup({ + state: { activeChains: [CHAIN_MAINNET] }, + }); + + triggerStatusChanged({ chainIds: [CHAIN_MAINNET], status: 'up' }); + + expect(onActiveChainsUpdated).not.toHaveBeenCalled(); + + cleanup(); + }); + + it('swallows errors thrown while handling the event', () => { + const onActiveChainsUpdated = jest.fn().mockImplementation(() => { + throw new Error('listener boom'); + }); + const { triggerStatusChanged, cleanup } = setup({ + onActiveChainsUpdated, + }); + + expect(() => + triggerStatusChanged({ chainIds: [CHAIN_MAINNET], status: 'up' }), + ).not.toThrow(); + + cleanup(); + }); + }); + + describe('destroy', () => { + it('does not throw and clears internal subscriptions', () => { + const { dataSource, cleanup } = setup(); + + expect(() => dataSource.destroy()).not.toThrow(); + + cleanup(); + }); + + it('can be called multiple times safely', () => { + const { dataSource, cleanup } = setup(); + + expect(() => { + dataSource.destroy(); + dataSource.destroy(); + }).not.toThrow(); + + cleanup(); + }); + }); + + describe('createAccountActivityDataSource', () => { + it('returns an AccountActivityDataSource instance', () => { + const rootMessenger: RootMessenger = new Messenger({ + namespace: MOCK_ANY_NAMESPACE, + }); + const assetsControllerMessenger: AssetsControllerMessenger = + new Messenger({ + namespace: 'AssetsController', + parent: rootMessenger, + }); + rootMessenger.delegate({ + messenger: assetsControllerMessenger, + actions: [], + events: [ + 'AccountActivityService:balanceUpdated', + 'AccountActivityService:statusChanged', + ], + }); + + const dataSource = createAccountActivityDataSource({ + messenger: assetsControllerMessenger, + onActiveChainsUpdated: jest.fn(), + getAssetType: () => 'native', + onAssetsUpdate: jest.fn(), + }); + + expect(dataSource).toBeInstanceOf(AccountActivityDataSource); + + dataSource.destroy(); + rootMessenger.clearSubscriptions(); + }); + }); +}); diff --git a/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts new file mode 100644 index 00000000000..3db0c632012 --- /dev/null +++ b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts @@ -0,0 +1,439 @@ +import type { + AccountActivityServiceBalanceUpdatedEvent, + AccountActivityServiceStatusChangedEvent, + BalanceUpdate, +} from '@metamask/core-backend'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; +import { isCaipChainId } from '@metamask/utils'; +import BigNumberJS from 'bignumber.js'; + +import type { AssetsControllerMessenger } from '../AssetsController.js'; +import { projectLogger, createModuleLogger } from '../logger.js'; +import type { + AssetBalance, + AssetMetadata, + ChainId, + Caip19AssetId, + DataRequest, + DataResponse, +} from '../types.js'; +import { AbstractDataSource } from './AbstractDataSource.js'; +import type { DataSourceState } from './AbstractDataSource.js'; + +// ============================================================================ +// CONSTANTS +// ============================================================================ + +const CONTROLLER_NAME = 'AccountActivityDataSource'; + +const log = createModuleLogger(projectLogger, CONTROLLER_NAME); + +// ============================================================================ +// BALANCE UPDATE PROCESSING +// ============================================================================ + +/** + * Convert AccountActivityMessage balance updates into a {@link DataResponse} + * for AssetsController. + * + * @param updates - Balance updates from account-activity websocket payload. + * @param accountId - Internal account UUID. + * @param getAssetType - Resolver for asset metadata type. + * @returns DataResponse with merge mode when balances are present. + */ +function processAccountActivityBalanceUpdates( + updates: BalanceUpdate[], + accountId: string, + getAssetType: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl', +): DataResponse { + const assetsBalance = Object.create(null) as Record< + string, + Record + >; + assetsBalance[accountId] = Object.create(null) as Record< + Caip19AssetId, + AssetBalance + >; + const assetsMetadata = Object.create(null) as Record< + Caip19AssetId, + AssetMetadata + >; + + for (const update of updates) { + const { asset, postBalance } = update; + + if (!asset || !postBalance) { + continue; + } + + const assetId = asset.type as Caip19AssetId; + + if (asset.decimals === undefined) { + continue; + } + + const rawBalanceStr = postBalance.amount.startsWith('0x') + ? BigInt(postBalance.amount).toString() + : postBalance.amount; + + const humanReadableAmount = new BigNumberJS(rawBalanceStr) + .dividedBy(new BigNumberJS(10).pow(asset.decimals)) + .toFixed(); + + assetsBalance[accountId][assetId] = { + amount: humanReadableAmount, + }; + + assetsMetadata[assetId] = { + type: getAssetType(assetId), + symbol: asset.unit, + name: asset.unit, + decimals: asset.decimals, + }; + } + + const response: DataResponse = { updateMode: 'merge' }; + if (Object.keys(assetsBalance[accountId]).length > 0) { + response.assetsBalance = assetsBalance; + response.assetsInfo = assetsMetadata; + } + + return response; +} + +// ============================================================================ +// MESSENGER TYPES +// ============================================================================ + +/** Allowed events that AccountActivityDataSource subscribes to. */ +export type AccountActivityDataSourceAllowedEvents = + | AccountActivityServiceBalanceUpdatedEvent + | AccountActivityServiceStatusChangedEvent; + +// ============================================================================ +// STATE +// ============================================================================ + +export type AccountActivityDataSourceState = DataSourceState; + +const defaultState: AccountActivityDataSourceState = { + activeChains: [], +}; + +// ============================================================================ +// OPTIONS +// ============================================================================ + +export type AccountActivityDataSourceOptions = { + /** The AssetsController messenger (shared by all data sources). */ + messenger: AssetsControllerMessenger; + /** Called when active chains are updated. Pass dataSourceName so the controller knows the source. */ + onActiveChainsUpdated: ( + dataSourceName: string, + chains: ChainId[], + previousChains: ChainId[], + ) => void; + /** Returns the asset type ('native' | 'erc20' | 'spl') for a given CAIP-19 asset ID. */ + getAssetType: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl'; + /** + * Pushes decoded balance updates to the controller (bound to + * `AssetsController.handleAssetsUpdate`). AADS is event-driven and never + * takes part in the controller's subscribe handoff, so it receives this + * callback directly instead of via a `SubscriptionRequest`. + */ + onAssetsUpdate: ( + response: DataResponse, + request?: DataRequest, + ) => void | Promise; + state?: Partial; +}; + +// ============================================================================ +// ACCOUNT ACTIVITY DATA SOURCE +// ============================================================================ + +/** + * Data source that consumes real-time updates from `AccountActivityService`. + * + * `AccountActivityService` owns the WebSocket connection and channel + * subscriptions; this data source is a thin consumer of the two high-level + * events that it publishes: + * + * - `AccountActivityService:balanceUpdated` — post-transaction balances for the + * subscribed account(s). The address is resolved against the wallet's + * selected account group, transformed into a {@link DataResponse} (via + * {@link processAccountActivityBalanceUpdates}), and pushed to the controller + * through the injected `onAssetsUpdate` callback. + * - `AccountActivityService:statusChanged` — per-chain "up"/"down" notifications. + * A chain reported "up" is claimed as an active chain (WebSocket is providing + * real-time data); a chain reported "down" is released so polling data sources + * can take over. Each change is applied directly via `updateActiveChains`, + * which notifies the controller through `onActiveChainsUpdated`. + * + * This data source does NOT debounce, jitter, or gate chain updates itself. + * Coalescing (collapsing bursts) and jitter (staggering the WS-subscribe herd + * across clients) live in `AssetsController`, where the expensive re-subscribe + * happens and where updates from all data sources converge. `activeChains` are + * never seeded from the accounts API; they only ever reflect live + * `statusChanged` events (the service flushes all tracked chains as "down" when + * the WebSocket disconnects, releasing them). + * + * It subscribes to these events in its constructor (the same way + * `TokenBalancesController` does). Unlike the polling data sources it does not + * take part in the controller's subscribe/unsubscribe handoff: `subscribe` is a + * no-op, balance updates are routed by resolving the address against the + * wallet, and updates are pushed through the injected `onAssetsUpdate` callback. + */ +export class AccountActivityDataSource extends AbstractDataSource< + typeof CONTROLLER_NAME, + AccountActivityDataSourceState +> { + readonly #messenger: AssetsControllerMessenger; + + readonly #onActiveChainsUpdated: ( + dataSourceName: string, + chains: ChainId[], + previousChains: ChainId[], + ) => void; + + readonly #getAssetType: ( + assetId: Caip19AssetId, + ) => 'native' | 'erc20' | 'spl'; + + readonly #onAssetsUpdate: ( + response: DataResponse, + request?: DataRequest, + ) => void | Promise; + + /** Unsubscribe handles for messenger event subscriptions. */ + readonly #eventUnsubscribes: (() => void)[] = []; + + constructor(options: AccountActivityDataSourceOptions) { + super(CONTROLLER_NAME, { + ...defaultState, + ...options.state, + }); + + this.#messenger = options.messenger; + this.#onActiveChainsUpdated = options.onActiveChainsUpdated; + this.#getAssetType = options.getAssetType; + this.#onAssetsUpdate = options.onAssetsUpdate; + + this.#subscribeToEvents(); + } + + // ============================================================================ + // EVENT SUBSCRIPTIONS + // ============================================================================ + + #subscribeToEvents(): void { + const unsubscribeBalance = this.#messenger.subscribe( + 'AccountActivityService:balanceUpdated', + (event) => this.#onBalanceUpdated(event), + ); + const unsubscribeStatus = this.#messenger.subscribe( + 'AccountActivityService:statusChanged', + this.#onAccountActivityStatusChanged, + ); + + if (typeof unsubscribeBalance === 'function') { + this.#eventUnsubscribes.push(unsubscribeBalance); + } + if (typeof unsubscribeStatus === 'function') { + this.#eventUnsubscribes.push(unsubscribeStatus); + } + } + + // ============================================================================ + // SUBSCRIBE / UNSUBSCRIBE + // ============================================================================ + + /** + * AADS is event-driven and chain-agnostic: it never participates in the + * controller's subscribe/unsubscribe handoff. Incoming `balanceUpdated` + * events are routed by resolving the address against the wallet's accounts + * (see `#onBalanceUpdated`), and updates are pushed through the injected + * `#onAssetsUpdate` callback. This override exists only to satisfy the + * abstract contract and is intentionally a no-op. + */ + async subscribe(): Promise { + // Intentionally empty — see method doc. + } + + // ============================================================================ + // BALANCE UPDATES + // ============================================================================ + + #onBalanceUpdated({ + address, + chain, + updates, + }: { + address: string; + chain: string; + updates: BalanceUpdate[]; + }): void { + try { + if (!address || !chain || !updates || updates.length === 0) { + return; + } + + const account = this.#findAccountForAddress(address); + if (!account) { + return; + } + + const chainId = chain as ChainId; + + const response = processAccountActivityBalanceUpdates( + updates, + account.id, + (assetId) => this.#getAssetType(assetId), + ); + + if (!response.assetsBalance) { + return; + } + + const request: DataRequest = { + accountsWithSupportedChains: [{ account, supportedChains: [chainId] }], + chainIds: [chainId], + dataTypes: ['balance', 'metadata'], + }; + + Promise.resolve(this.#onAssetsUpdate(response, request)).catch( + (error) => { + log('Failed to report balance update', { error }); + }, + ); + } catch (error) { + log('Error handling balance update', error); + } + } + + /** + * Find the wallet account matching the given activity address. EVM addresses + * are matched case-insensitively; other namespaces are matched exactly. + * + * The candidate set is the accounts of the selected account group, resolved + * from the wallet at event time (rather than from a stored subscription), so + * AADS does not need to be re-subscribed whenever the account set changes. + * + * @param address - The account address from the activity message. + * @returns The matching account, or null. + */ + #findAccountForAddress(address: string): InternalAccount | null { + const isEvm = address.startsWith('0x'); + + return ( + this.#getWalletAccounts().find((candidate) => + isEvm + ? candidate.address.toLowerCase() === address.toLowerCase() + : candidate.address === address, + ) ?? null + ); + } + + /** + * Resolve the accounts of the selected account group from the wallet, + * mirroring `AssetsController.#getSelectedAccounts`. + * + * @returns The accounts of the selected account group. + */ + #getWalletAccounts(): InternalAccount[] { + const accounts = this.#messenger.call( + 'AccountTreeController:getAccountsFromSelectedAccountGroup', + ); + if (accounts.length > 0) { + return accounts; + } + + const selectedAccount = this.#messenger.call( + 'AccountsController:getSelectedAccount', + ); + return selectedAccount ? [selectedAccount] : []; + } + + // ============================================================================ + // STATUS CHANGES + // ============================================================================ + + /** + * Handle a `statusChanged` notification. Chains reported "up" are claimed as + * active (WebSocket provides real-time data); chains reported "down" are + * released so polling data sources take over. The change is applied directly; + * coalescing and jitter are handled by `AssetsController`. + * + * @param options - The status change notification. + * @param options.chainIds - The CAIP-2 chain IDs whose status changed. + * @param options.status - Whether the chains are `up` or `down`. + * @param options.timestamp - Optional timestamp of the status change. + */ + readonly #onAccountActivityStatusChanged = ({ + chainIds, + status, + }: { + chainIds: string[]; + status: 'up' | 'down'; + timestamp?: number; + }): void => { + try { + // Act on every namespace (eip155, solana, etc.); AssetsController is + // multichain. Only skip identifiers that are not valid CAIP-2 chain IDs. + const validChains = chainIds.filter((chainId) => + isCaipChainId(chainId), + ) as ChainId[]; + + if (validChains.length === 0) { + return; + } + + const next = new Set(this.state.activeChains); + if (status === 'up') { + for (const chainId of validChains) { + next.add(chainId); + } + } else { + for (const chainId of validChains) { + next.delete(chainId); + } + } + + const previous = [...this.state.activeChains]; + this.updateActiveChains(Array.from(next), (updatedChains) => + this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), + ); + } catch (error) { + log('Error handling status change', error); + } + }; + + // ============================================================================ + // CLEANUP + // ============================================================================ + + destroy(): void { + for (const unsubscribe of this.#eventUnsubscribes) { + unsubscribe(); + } + this.#eventUnsubscribes.length = 0; + + super.destroy(); + } +} + +// ============================================================================ +// FACTORY FUNCTION +// ============================================================================ + +/** + * Creates an AccountActivityDataSource instance. + * + * @param options - Configuration options for the data source. + * @returns A new AccountActivityDataSource instance. + */ +export function createAccountActivityDataSource( + options: AccountActivityDataSourceOptions, +): AccountActivityDataSource { + return new AccountActivityDataSource(options); +} diff --git a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts deleted file mode 100644 index 5f6887deba5..00000000000 --- a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts +++ /dev/null @@ -1,1417 +0,0 @@ -/* eslint-disable jest/unbound-method */ -import type { - ApiPlatformClient, - ServerNotificationMessage, - WebSocketSubscription, -} from '@metamask/core-backend'; -import { WebSocketState } from '@metamask/core-backend'; -import type { InternalAccount } from '@metamask/keyring-internal-api'; -import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; -import type { MockAnyNamespace } from '@metamask/messenger'; - -import type { AssetsControllerMessenger } from '../AssetsController.js'; -import type { Caip19AssetId, ChainId, DataRequest } from '../types.js'; -import { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, -} from './BackendWebsocketDataSource.js'; -import type { - BackendWebsocketDataSourceAllowedActions, - BackendWebsocketDataSourceAllowedEvents, -} from './BackendWebsocketDataSource.js'; - -type AllActions = BackendWebsocketDataSourceAllowedActions; -type AllEvents = BackendWebsocketDataSourceAllowedEvents; -type RootMessenger = Messenger; - -const CHAIN_MAINNET = 'eip155:1' as ChainId; -const CHAIN_POLYGON = 'eip155:137' as ChainId; -const CHAIN_BASE = 'eip155:8453' as ChainId; -const MOCK_ADDRESS = '0x1234567890123456789012345678901234567890'; - -type SetupResult = { - controller: BackendWebsocketDataSource; - messenger: RootMessenger; - wsSubscribeMock: jest.Mock; - getConnectionInfoMock: jest.Mock; - findSubscriptionsMock: jest.Mock; - addChannelCallbackMock: jest.Mock; - removeChannelCallbackMock: jest.Mock; - assetsUpdateHandler: jest.Mock; - activeChainsUpdateHandler: jest.Mock; - triggerConnectionStateChange: (state: WebSocketState) => void; - triggerActiveChainsUpdate: (chains: ChainId[]) => void; -}; - -function createMockAccount( - overrides?: Partial, -): InternalAccount { - return { - id: 'mock-account-id', - address: MOCK_ADDRESS, - options: {}, - methods: [], - type: 'eip155:eoa', - scopes: ['eip155:0'], - metadata: { - name: 'Test Account', - keyring: { type: 'HD Key Tree' }, - importTime: Date.now(), - lastSelected: Date.now(), - }, - ...overrides, - } as InternalAccount; -} - -function createDataRequest( - overrides?: Partial & { accounts?: InternalAccount[] }, -): DataRequest { - const chainIds = overrides?.chainIds ?? [CHAIN_MAINNET]; - const accounts = overrides?.accounts ?? [createMockAccount()]; - const { accounts: _a, ...rest } = overrides ?? {}; - return { - chainIds, - accountsWithSupportedChains: accounts.map((a) => ({ - account: a, - supportedChains: chainIds, - })), - dataTypes: ['balance'], - ...rest, - }; -} - -function createMockWsSubscription( - channels: string[] = [], -): WebSocketSubscription { - return { - unsubscribe: jest.fn().mockResolvedValue(undefined), - channels, - } as unknown as WebSocketSubscription; -} - -function createMockNotification( - overrides: Partial & { - data: Record; - }, -): ServerNotificationMessage { - return { - event: 'notification', - channel: 'test-channel', - timestamp: Date.now(), - ...overrides, - }; -} - -function setupController( - options: { - initialActiveChains?: ChainId[]; - connectionState?: WebSocketState; - } = {}, -): SetupResult { - const { - initialActiveChains = [], - connectionState = WebSocketState.CONNECTED, - } = options; - - const rootMessenger = new Messenger({ - namespace: MOCK_ANY_NAMESPACE, - }); - - const controllerMessenger = new Messenger< - 'BackendWebsocketDataSource', - AllActions, - AllEvents, - RootMessenger - >({ - namespace: 'BackendWebsocketDataSource', - parent: rootMessenger, - }); - - rootMessenger.delegate({ - messenger: controllerMessenger, - actions: [ - 'BackendWebSocketService:subscribe', - 'BackendWebSocketService:getConnectionInfo', - 'BackendWebSocketService:findSubscriptionsByChannelPrefix', - 'BackendWebSocketService:addChannelCallback', - 'BackendWebSocketService:removeChannelCallback', - ], - events: ['BackendWebSocketService:connectionStateChanged'], - }); - - const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined); - const activeChainsUpdateHandler = jest.fn(); - const wsSubscribeMock = jest - .fn() - .mockResolvedValue(createMockWsSubscription()); - const addChannelCallbackMock = jest.fn(); - const removeChannelCallbackMock = jest.fn().mockReturnValue(true); - const getConnectionInfoMock = jest.fn().mockReturnValue({ - state: connectionState, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }); - const findSubscriptionsMock = jest.fn().mockReturnValue([]); - - rootMessenger.registerActionHandler( - 'BackendWebSocketService:subscribe', - wsSubscribeMock, - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:getConnectionInfo', - getConnectionInfoMock, - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:findSubscriptionsByChannelPrefix', - findSubscriptionsMock, - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:addChannelCallback', - addChannelCallbackMock, - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:removeChannelCallback', - removeChannelCallbackMock, - ); - - const queryApiClient = { - accounts: { - fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ - fullSupport: initialActiveChains.map((chainId) => { - const [, ref] = chainId.split(':'); - return parseInt(ref, 10); - }), - }), - }, - }; - - const getAssetTypeFn = ( - assetId: Caip19AssetId, - ): 'native' | 'erc20' | 'spl' => { - if (assetId.includes('/slip44:')) { - return 'native'; - } - if (assetId.startsWith('solana:') && assetId.includes('/token:')) { - return 'spl'; - } - return 'erc20'; - }; - - const controller = new BackendWebsocketDataSource({ - messenger: controllerMessenger as unknown as AssetsControllerMessenger, - queryApiClient: queryApiClient as unknown as ApiPlatformClient, - onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => - activeChainsUpdateHandler(dataSourceName, chains, previousChains), - getAssetType: getAssetTypeFn, - state: { activeChains: initialActiveChains }, - }); - - const triggerConnectionStateChange = (state: WebSocketState): void => { - rootMessenger.publish('BackendWebSocketService:connectionStateChanged', { - state, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }); - }; - - const triggerActiveChainsUpdate = (chains: ChainId[]): void => { - controller.setActiveChainsFromAccountsApi(chains); - activeChainsUpdateHandler( - 'BackendWebsocketDataSource', - chains, - initialActiveChains, - ); - }; - - return { - controller, - messenger: rootMessenger, - wsSubscribeMock, - getConnectionInfoMock, - findSubscriptionsMock, - addChannelCallbackMock, - removeChannelCallbackMock, - assetsUpdateHandler, - activeChainsUpdateHandler, - triggerConnectionStateChange, - triggerActiveChainsUpdate, - }; -} - -describe('BackendWebsocketDataSource', () => { - afterEach(() => { - jest.clearAllMocks(); - }); - - it('initializes with correct name', () => { - const { controller } = setupController(); - expect(controller.getName()).toBe('BackendWebsocketDataSource'); - controller.destroy(); - }); - - it('exposes getActiveChains on instance', async () => { - const { controller } = setupController(); - - const chains = await controller.getActiveChains(); - expect(chains).toStrictEqual([]); - - controller.destroy(); - }); - - it('updates active chains when AccountsApiDataSource publishes update', async () => { - const { controller, triggerActiveChainsUpdate, activeChainsUpdateHandler } = - setupController(); - - triggerActiveChainsUpdate([CHAIN_MAINNET, CHAIN_POLYGON]); - - const chains = await controller.getActiveChains(); - expect(chains).toStrictEqual([CHAIN_MAINNET, CHAIN_POLYGON]); - expect(activeChainsUpdateHandler).toHaveBeenCalledWith( - 'BackendWebsocketDataSource', - [CHAIN_MAINNET, CHAIN_POLYGON], - [], - ); - - controller.destroy(); - }); - - it('updateSupportedChains updates active chains', async () => { - const { controller, activeChainsUpdateHandler } = setupController(); - - controller.updateSupportedChains([CHAIN_MAINNET, CHAIN_BASE]); - - const chains = await controller.getActiveChains(); - expect(chains).toStrictEqual([CHAIN_MAINNET, CHAIN_BASE]); - expect(activeChainsUpdateHandler).toHaveBeenCalledWith( - 'BackendWebsocketDataSource', - [CHAIN_MAINNET, CHAIN_BASE], - [], - ); - - controller.destroy(); - }); - - it('subscribe creates eip155 channel when no request chains match (eip155 account only)', async () => { - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_POLYGON] }), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledWith( - expect.objectContaining({ - channels: [ - `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - ], - channelType: 'account-activity.v1', - callback: expect.any(Function), - }), - ); - - controller.destroy(); - }); - - it('subscribe creates WebSocket subscription when connected', async () => { - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - // EIP-155 account only -> eip155 channel with lowercase hex - expect(wsSubscribeMock).toHaveBeenCalledWith( - expect.objectContaining({ - channels: [ - `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - ], - channelType: 'account-activity.v1', - callback: expect.any(Function), - }), - ); - - controller.destroy(); - }); - - it('subscribe stores pending subscription when disconnected', async () => { - const { - controller, - wsSubscribeMock, - getConnectionInfoMock, - triggerConnectionStateChange, - } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.DISCONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).not.toHaveBeenCalled(); - - getConnectionInfoMock.mockReturnValue({ - state: WebSocketState.CONNECTED, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }); - - triggerConnectionStateChange(WebSocketState.CONNECTED); - await new Promise(process.nextTick); - - // Stale pending subscriptions are cleared on reconnect rather than - // being re-processed. The chain reclaim via updateActiveChains - // triggers onActiveChainsUpdated, causing AssetsController to create - // fresh subscriptions with current data. - expect(wsSubscribeMock).not.toHaveBeenCalled(); - - controller.destroy(); - }); - - it('subscribe creates channels for multiple namespaces with correct address format per namespace', async () => { - const solanaAddress = '7xKXtg2CW87d97TXJSDpbD5jBkheTqA83TZRuJosgAsU'; - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [ - CHAIN_MAINNET, - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId, - ], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ - chainIds: [ - CHAIN_MAINNET, - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId, - ], - accountsWithSupportedChains: [ - { - account: createMockAccount(), - supportedChains: [CHAIN_MAINNET], - }, - { - account: createMockAccount({ - id: 'solana-account-id', - address: solanaAddress, - type: 'solana:data-account', - scopes: ['solana:0'], - }), - supportedChains: [ - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId, - ], - }, - ], - }), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - // EIP-155: lowercase hex; Solana: base58 as-is - expect(wsSubscribeMock).toHaveBeenCalledWith( - expect.objectContaining({ - channels: expect.arrayContaining([ - `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - `account-activity.v1.solana:0:${solanaAddress}`, - ]), - }), - ); - - controller.destroy(); - }); - - it('subscribe update only changes chains if addresses unchanged', async () => { - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET, CHAIN_POLYGON], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_MAINNET] }), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON] }), - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - - controller.destroy(); - }); - - it('subscribe update re-subscribes when addresses change', async () => { - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - - const newAddress = '0xabcdef1234567890abcdef1234567890abcdef12'; - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ - accountsWithSupportedChains: [ - { - account: createMockAccount({ address: newAddress }), - supportedChains: [CHAIN_MAINNET], - }, - ], - chainIds: [CHAIN_MAINNET], - }), - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(2); - - controller.destroy(); - }); - - it('subscribe update treats checksummed and lowercase EVM addresses as unchanged', async () => { - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ - accountsWithSupportedChains: [ - { - account: createMockAccount({ - address: `0x${MOCK_ADDRESS.slice(2).toUpperCase()}`, - }), - supportedChains: [CHAIN_MAINNET], - }, - ], - chainIds: [CHAIN_MAINNET], - }), - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - - controller.destroy(); - }); - - it('serializes concurrent subscribe calls so the last address wins', async () => { - const addressA = MOCK_ADDRESS; - const addressB = '0xabcdef1234567890abcdef1234567890abcdef12'; - let resolveFirstSubscribe: (() => void) | undefined; - const firstSubscribeGate = new Promise((resolve) => { - resolveFirstSubscribe = resolve; - }); - - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - wsSubscribeMock - .mockImplementationOnce(async () => { - await firstSubscribeGate; - return createMockWsSubscription([ - `account-activity.v1.eip155:0:${addressA.toLowerCase()}`, - ]); - }) - .mockResolvedValue( - createMockWsSubscription([ - `account-activity.v1.eip155:0:${addressB.toLowerCase()}`, - ]), - ); - - const firstSubscribe = controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ - accountsWithSupportedChains: [ - { - account: createMockAccount({ address: addressA }), - supportedChains: [CHAIN_MAINNET], - }, - ], - }), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - const secondSubscribe = controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ - accountsWithSupportedChains: [ - { - account: createMockAccount({ address: addressB }), - supportedChains: [CHAIN_MAINNET], - }, - ], - }), - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - - await new Promise(process.nextTick); - resolveFirstSubscribe?.(); - await Promise.all([firstSubscribe, secondSubscribe]); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(2); - expect(wsSubscribeMock.mock.calls[1][0].channels).toStrictEqual([ - `account-activity.v1.eip155:0:${addressB.toLowerCase()}`, - ]); - - controller.destroy(); - }); - - it('unsubscribe cleans up WebSocket subscription', async () => { - const channel = `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`; - const mockWsSubscription = createMockWsSubscription([channel]); - const { controller, wsSubscribeMock, removeChannelCallbackMock } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - wsSubscribeMock.mockResolvedValueOnce(mockWsSubscription); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - await controller.unsubscribe('sub-1'); - - expect(mockWsSubscription.unsubscribe).toHaveBeenCalled(); - expect(removeChannelCallbackMock).toHaveBeenCalledWith(channel); - - controller.destroy(); - }); - - it('registers channel callbacks as fallback when subscriptionId does not match', async () => { - const channel = `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`; - const mockWsSubscription = createMockWsSubscription([channel]); - const onAssetsUpdate = jest.fn().mockResolvedValue(undefined); - const { controller, wsSubscribeMock, addChannelCallbackMock } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - wsSubscribeMock.mockResolvedValueOnce(mockWsSubscription); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate, - }); - - expect(addChannelCallbackMock).toHaveBeenCalledWith( - expect.objectContaining({ channelName: channel }), - ); - - const channelCallback = addChannelCallbackMock.mock.calls.find( - ([args]) => args.channelName === channel, - )?.[0].callback; - - expect(channelCallback).toBeDefined(); - - channelCallback( - createMockNotification({ - channel: `account-activity.v1.eip155:42161:${MOCK_ADDRESS.toLowerCase()}`, - subscriptionId: 'stale-server-sub-id', - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', - decimals: 6, - }, - postBalance: { amount: '1000000' }, - }, - ], - }, - }), - ); - - await new Promise(process.nextTick); - - expect(onAssetsUpdate).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.any(Object), - }), - }), - expect.objectContaining({ - accountsWithSupportedChains: expect.any(Array), - }), - ); - - controller.destroy(); - }); - - it('still stores subscription state when channel callback registration fails', async () => { - const channel = `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`; - const onAssetsUpdate = jest.fn().mockResolvedValue(undefined); - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - const rootMessenger = new Messenger< - MockAnyNamespace, - AllActions, - AllEvents - >({ namespace: MOCK_ANY_NAMESPACE }); - const controllerMessenger = new Messenger< - 'BackendWebsocketDataSource', - AllActions, - AllEvents, - RootMessenger - >({ - namespace: 'BackendWebsocketDataSource', - parent: rootMessenger, - }); - - rootMessenger.delegate({ - messenger: controllerMessenger, - actions: [ - 'BackendWebSocketService:subscribe', - 'BackendWebSocketService:getConnectionInfo', - 'BackendWebSocketService:addChannelCallback', - ], - events: ['BackendWebSocketService:connectionStateChanged'], - }); - - rootMessenger.registerActionHandler( - 'BackendWebSocketService:subscribe', - ({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription([channel])); - }, - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:getConnectionInfo', - () => ({ - state: WebSocketState.CONNECTED, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }), - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:addChannelCallback', - () => { - throw new Error( - 'A handler for BackendWebSocketService:addChannelCallback has not been delegated to AssetsController', - ); - }, - ); - - const controller = new BackendWebsocketDataSource({ - messenger: controllerMessenger as unknown as AssetsControllerMessenger, - queryApiClient: { - accounts: { - fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ - fullSupport: [1], - }), - }, - } as unknown as ApiPlatformClient, - onActiveChainsUpdated: jest.fn(), - getAssetType: (): 'erc20' => 'erc20', - state: { activeChains: [CHAIN_MAINNET] }, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate, - }); - - notificationCallback( - createMockNotification({ - channel, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', - decimals: 6, - }, - postBalance: { amount: '1000000' }, - }, - ], - }, - }), - ); - - await new Promise(process.nextTick); - - expect(onAssetsUpdate).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.any(Object), - }), - }), - expect.objectContaining({ dataTypes: ['balance'] }), - ); - - controller.destroy(); - }); - - it('handles WebSocket disconnect by releasing chains and reclaiming on reconnect', async () => { - const { - controller, - wsSubscribeMock, - getConnectionInfoMock, - activeChainsUpdateHandler, - triggerConnectionStateChange, - } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - - getConnectionInfoMock.mockReturnValue({ - state: WebSocketState.DISCONNECTED, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }); - - triggerConnectionStateChange(WebSocketState.DISCONNECTED); - - getConnectionInfoMock.mockReturnValue({ - state: WebSocketState.CONNECTED, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }); - - activeChainsUpdateHandler.mockClear(); - triggerConnectionStateChange(WebSocketState.CONNECTED); - await new Promise(process.nextTick); - - // Stale pending subscriptions are NOT re-processed on reconnect. - // Instead, chain reclaim fires onActiveChainsUpdated so - // AssetsController creates fresh subscriptions with current data. - expect(wsSubscribeMock).toHaveBeenCalledTimes(1); - expect(activeChainsUpdateHandler).toHaveBeenCalledWith( - 'BackendWebsocketDataSource', - [CHAIN_MAINNET], - expect.any(Array), - ); - - controller.destroy(); - }); - - it('processes balance update notification correctly', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_BASE], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_BASE] }), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - const notification = createMockNotification({ - channel: `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_BASE }, - updates: [ - { - asset: { - type: 'eip155:8453/slip44:60', - unit: 'ETH', - decimals: 18, - }, - postBalance: { - amount: '0x8ac7230489e80000', - }, - }, - ], - }, - }); - - notificationCallback(notification); - await new Promise(process.nextTick); - - // Raw 10e18 wei (0x8ac7230489e80000) with 18 decimals → human-readable "10" - expect(assetsUpdateHandler).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.objectContaining({ - 'eip155:8453/slip44:60': { amount: '10' }, - }), - }), - assetsInfo: expect.objectContaining({ - 'eip155:8453/slip44:60': expect.objectContaining({ - type: 'native', - symbol: 'ETH', - decimals: 18, - }), - }), - }), - expect.objectContaining({ - dataTypes: ['balance'], - accountsWithSupportedChains: expect.any(Array), - }), - ); - - controller.destroy(); - }); - - it('processes ERC20 token balance update', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - const notification = createMockNotification({ - channel: `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', - unit: 'USDC', - decimals: 6, - }, - postBalance: { - amount: '1000000', - }, - }, - ], - }, - }); - - notificationCallback(notification); - await new Promise(process.nextTick); - - // Raw 1000000 (1 USDC) with 6 decimals → human-readable "1" - expect(assetsUpdateHandler).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.objectContaining({ - 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48': { - amount: '1', - }, - }), - }), - assetsInfo: expect.objectContaining({ - 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48': - expect.objectContaining({ - type: 'erc20', - symbol: 'USDC', - decimals: 6, - }), - }), - }), - expect.objectContaining({ - dataTypes: ['balance'], - accountsWithSupportedChains: expect.any(Array), - }), - ); - - controller.destroy(); - }); - - it('converts raw WebSocket balance (hex) to human-readable using asset decimals', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - // 0x26f0e5 = 2552037 raw; USDC 6 decimals → 2.552037 - const notification = createMockNotification({ - channel: `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', - unit: 'USDC', - decimals: 6, - }, - postBalance: { - amount: '0x26f0e5', - }, - }, - ], - }, - }); - - notificationCallback(notification); - await new Promise(process.nextTick); - - // assetId key is as in notification (mixed case) - expect(assetsUpdateHandler).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.objectContaining({ - 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48': { - amount: '2.552037', - }, - }), - }), - }), - expect.objectContaining({ - dataTypes: ['balance'], - accountsWithSupportedChains: expect.any(Array), - }), - ); - - controller.destroy(); - }); - - it('emits plain decimal (not exponent form) for sub-1e-7 dust balances', async () => { - // Regression for MMBUGS-772: BigNumber's default EXPONENTIAL_AT makes - // `.toString()` emit "1e-18" for tiny values, which crashes downstream - // BigInt() consumers. The source uses `.toFixed()` to stay in plain form. - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - // 1 wei of an 18-decimal token = 1e-18 — squarely past BigNumber's - // default exponential threshold. - const notification = createMockNotification({ - channel: `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', - unit: 'TEST', - decimals: 18, - }, - postBalance: { - amount: '0x1', - }, - }, - ], - }, - }); - - notificationCallback(notification); - await new Promise(process.nextTick); - - expect(assetsUpdateHandler).toHaveBeenCalledWith( - expect.objectContaining({ - assetsBalance: expect.objectContaining({ - 'mock-account-id': expect.objectContaining({ - 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48': { - amount: '0.000000000000000001', - }, - }), - }), - }), - expect.objectContaining({ - dataTypes: ['balance'], - accountsWithSupportedChains: expect.any(Array), - }), - ); - - controller.destroy(); - }); - - it('skips balance update when asset.decimals is missing', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - // No decimals on asset → update is skipped (we assume decimals are always present) - const notification = createMockNotification({ - channel: `account-activity.v1.eip155:0:${MOCK_ADDRESS.toLowerCase()}`, - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { - asset: { - type: 'eip155:1/erc20:0x0000000000000000000000000000000000000001', - unit: 'UNKNOWN', - decimals: undefined, - }, - postBalance: { - amount: '1000000000000000000', - }, - }, - ], - }, - }); - - notificationCallback(notification); - await new Promise(process.nextTick); - - expect(assetsUpdateHandler).not.toHaveBeenCalled(); - - controller.destroy(); - }); - - it('ignores notification with missing data', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - notificationCallback( - createMockNotification({ - data: { address: null, tx: null, updates: null }, - }), - ); - - expect(assetsUpdateHandler).not.toHaveBeenCalled(); - - controller.destroy(); - }); - - it('ignores notification for unknown account', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - notificationCallback( - createMockNotification({ - data: { - address: '0xunknown', - tx: { chain: CHAIN_MAINNET }, - updates: [ - { asset: { type: 'test' }, postBalance: { amount: '100' } }, - ], - }, - }), - ); - - expect(assetsUpdateHandler).not.toHaveBeenCalled(); - - controller.destroy(); - }); - - it('skips updates with missing asset or postBalance', async () => { - const { controller, wsSubscribeMock, assetsUpdateHandler } = - setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - let notificationCallback: ( - notification: ServerNotificationMessage, - ) => void = () => undefined; - - wsSubscribeMock.mockImplementation(({ callback }) => { - notificationCallback = callback; - return Promise.resolve(createMockWsSubscription()); - }); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - notificationCallback( - createMockNotification({ - data: { - address: MOCK_ADDRESS, - tx: { chain: CHAIN_MAINNET }, - updates: [ - { asset: null, postBalance: { amount: '100' } }, - { asset: { type: 'test' }, postBalance: null }, - ], - }, - }), - ); - - expect(assetsUpdateHandler).not.toHaveBeenCalled(); - - controller.destroy(); - }); - - it('destroy cleans up WebSocket subscriptions', async () => { - const mockWsSubscription = createMockWsSubscription(); - const { controller, wsSubscribeMock } = setupController({ - initialActiveChains: [CHAIN_MAINNET], - connectionState: WebSocketState.CONNECTED, - }); - - wsSubscribeMock.mockResolvedValueOnce(mockWsSubscription); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - - controller.destroy(); - - expect(mockWsSubscription.unsubscribe).toHaveBeenCalled(); - }); - - it('createBackendWebsocketDataSource factory creates instance', () => { - const rootMessenger = new Messenger< - MockAnyNamespace, - AllActions, - AllEvents - >({ - namespace: MOCK_ANY_NAMESPACE, - }); - - const controllerMessenger = new Messenger< - 'BackendWebsocketDataSource', - AllActions, - AllEvents, - RootMessenger - >({ - namespace: 'BackendWebsocketDataSource', - parent: rootMessenger, - }); - - rootMessenger.delegate({ - messenger: controllerMessenger, - actions: [ - 'BackendWebSocketService:subscribe', - 'BackendWebSocketService:getConnectionInfo', - 'BackendWebSocketService:findSubscriptionsByChannelPrefix', - ], - events: ['BackendWebSocketService:connectionStateChanged'], - }); - - rootMessenger.registerActionHandler( - 'BackendWebSocketService:subscribe', - jest.fn(), - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:getConnectionInfo', - jest.fn().mockReturnValue({ - state: WebSocketState.DISCONNECTED, - url: 'wss://test.example.com', - reconnectAttempts: 0, - timeout: 30000, - reconnectDelay: 1000, - maxReconnectDelay: 30000, - requestTimeout: 30000, - }), - ); - rootMessenger.registerActionHandler( - 'BackendWebSocketService:findSubscriptionsByChannelPrefix', - jest.fn(), - ); - - const queryApiClient = { - accounts: { - fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ - fullSupport: [], - }), - }, - }; - - const instance = createBackendWebsocketDataSource({ - messenger: controllerMessenger as unknown as AssetsControllerMessenger, - queryApiClient: queryApiClient as unknown as ApiPlatformClient, - onActiveChainsUpdated: jest.fn(), - }); - - expect(instance).toBeInstanceOf(BackendWebsocketDataSource); - expect(instance.getName()).toBe('BackendWebsocketDataSource'); - - instance.destroy(); - }); -}); diff --git a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts deleted file mode 100644 index e15290bedb2..00000000000 --- a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts +++ /dev/null @@ -1,843 +0,0 @@ -import type { - BackendWebSocketServiceActions, - BackendWebSocketServiceEvents, - ServerNotificationMessage, - WebSocketSubscription, - WebSocketState, - AccountActivityMessage, - BalanceUpdate, -} from '@metamask/core-backend'; -import type { ApiPlatformClient } from '@metamask/core-backend'; -import { - isCaipChainId, - KnownCaipNamespace, - toCaipChainId, -} from '@metamask/utils'; - -import type { AssetsControllerMessenger } from '../AssetsController.js'; -import { projectLogger, createModuleLogger } from '../logger.js'; -import type { ChainId, Caip19AssetId, DataResponse } from '../types.js'; -import { processAccountActivityBalanceUpdates } from '../utils/processAccountActivityBalanceUpdates.js'; -import { AbstractDataSource } from './AbstractDataSource.js'; -import type { - DataSourceState, - SubscriptionRequest, -} from './AbstractDataSource.js'; - -// ============================================================================ -// CONSTANTS -// ============================================================================ - -const CONTROLLER_NAME = 'BackendWebsocketDataSource'; -const CHANNEL_TYPE = 'account-activity.v1'; - -const log = createModuleLogger(projectLogger, CONTROLLER_NAME); - -// ============================================================================ -// MESSENGER TYPES -// ============================================================================ - -// Allowed actions that BackendWebsocketDataSource can call -export type BackendWebsocketDataSourceAllowedActions = - BackendWebSocketServiceActions; - -// Allowed events that BackendWebsocketDataSource can subscribe to -export type BackendWebsocketDataSourceAllowedEvents = - BackendWebSocketServiceEvents; - -// ============================================================================ -// STATE -// ============================================================================ - -export type BackendWebsocketDataSourceState = DataSourceState; - -const defaultState: BackendWebsocketDataSourceState = { - activeChains: [], -}; - -// ============================================================================ -// OPTIONS -// ============================================================================ - -export type BackendWebsocketDataSourceOptions = { - /** The AssetsController messenger (shared by all data sources). */ - messenger: AssetsControllerMessenger; - /** ApiPlatformClient for fetching supported networks at init (same as AccountsApiDataSource). */ - queryApiClient: ApiPlatformClient; - /** Called when active chains are updated. Pass dataSourceName so the controller knows the source. */ - onActiveChainsUpdated: ( - dataSourceName: string, - chains: ChainId[], - previousChains: ChainId[], - ) => void; - /** Returns the asset type ('native' | 'erc20' | 'spl') for a given CAIP-19 asset ID. */ - getAssetType: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl'; - state?: Partial; -}; - -// ============================================================================ -// HELPER FUNCTIONS -// ============================================================================ - -/** - * Extract namespace from a CAIP-2 chain ID. - * E.g., "eip155:1" -> "eip155", "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp" -> "solana" - * - * @param chainId - The CAIP-2 chain ID to extract namespace from. - * @returns The namespace portion of the chain ID. - */ -function extractNamespace(chainId: ChainId): string { - const [namespace] = chainId.split(':'); - return namespace; -} - -/** Namespaces we always subscribe to for account activity (EVM + Solana). */ -const ACCOUNT_ACTIVITY_NAMESPACES = ['eip155', 'solana'] as const; - -/** - * Get unique namespaces for account-activity subscriptions. - * Always includes eip155 and solana so we subscribe to both EVM and Solana account activity, - * plus any additional namespaces from the requested chain IDs. - * - * @param chainIds - Array of CAIP-2 chain IDs (from the subscription request). - * @returns Array of unique namespaces (at least eip155 and solana). - */ -function getNamespacesForAccountActivity(chainIds: ChainId[]): string[] { - const namespaces = new Set(ACCOUNT_ACTIVITY_NAMESPACES); - for (const chainId of chainIds) { - namespaces.add(extractNamespace(chainId)); - } - return Array.from(namespaces); -} - -/** - * Returns the address to use for account-activity subscription in the given namespace. - * EIP-155 accounts use hex (0x...) address; Solana accounts use base58. - * Returns null if this account type does not have an address in that namespace. - * - * @param account - Internal account (type + address). - * @param account.type - Account type (e.g. "eip155:eoa", "solana:data-account"). - * @param account.address - Account address (hex for eip155, base58 for solana). - * @param namespace - The chain namespace (e.g., "eip155", "solana"). - * @returns The address for that namespace, or null if the account does not support the namespace. - */ -function getAddressForAccountActivity( - account: { type: string; address: string }, - namespace: string, -): string | null { - if (namespace === 'eip155') { - return account.type.startsWith('eip155') ? account.address : null; - } - if (namespace === 'solana') { - return account.type.startsWith('solana') ? account.address : null; - } - // Other namespaces (e.g. from chainIds): use address if account type matches namespace - const typePrefix = `${namespace}:`; - return account.type.startsWith(typePrefix) ? account.address : null; -} - -/** - * Build WebSocket channel name for account activity using CAIP-10 wildcard format. - * Uses 0 as the chain reference to subscribe to all chains in the namespace. - * EIP-155 addresses are lowercased (hex); Solana addresses are left as-is (base58). - * - * @param namespace - The chain namespace (e.g., "eip155", "solana"). - * @param address - The account address (hex for eip155, base58 for solana). - * @returns The WebSocket channel name. - */ -function buildAccountActivityChannel( - namespace: string, - address: string, -): string { - const formatted = namespace === 'eip155' ? address.toLowerCase() : address; - return `${CHANNEL_TYPE}.${namespace}:0:${formatted}`; -} - -/** - * Normalize addresses for stable comparison when detecting account changes. - * - * @param address - Account address (hex or base58). - * @returns Normalized address for comparison. - */ -function normalizeAddressForComparison(address: string): string { - return address.startsWith('0x') ? address.toLowerCase() : address; -} - -/** - * Check whether subscribed account addresses changed (case-insensitive for EVM). - * - * @param nextAddresses - Addresses from the incoming subscribe request. - * @param existingAddresses - Addresses from the active subscription. - * @returns True when the address sets differ. - */ -function haveAddressesChanged( - nextAddresses: string[], - existingAddresses: string[], -): boolean { - if (nextAddresses.length !== existingAddresses.length) { - return true; - } - - const normalizedNext = nextAddresses - .map(normalizeAddressForComparison) - .sort(); - const normalizedExisting = existingAddresses - .map(normalizeAddressForComparison) - .sort(); - - return normalizedNext.some( - (address, index) => address !== normalizedExisting[index], - ); -} - -/** - * Normalize API chain identifier to CAIP-2 ChainId. - * Passes through strings already in CAIP-2 form (e.g. eip155:1, solana:5eykt...). - * Converts bare decimals to eip155:decimal. - * Uses @metamask/utils for CAIP parsing. - * - * @param chainIdOrDecimal - Chain ID string (CAIP-2 or decimal) or decimal number. - * @returns CAIP-2 ChainId. - */ -function toChainId(chainIdOrDecimal: number | string): ChainId { - if (typeof chainIdOrDecimal === 'string') { - if (isCaipChainId(chainIdOrDecimal)) { - return chainIdOrDecimal; - } - return toCaipChainId(KnownCaipNamespace.Eip155, chainIdOrDecimal); - } - return toCaipChainId(KnownCaipNamespace.Eip155, String(chainIdOrDecimal)); -} - -// Note: AccountActivityMessage and BalanceUpdate types are imported from @metamask/core-backend - -// ============================================================================ -// BACKEND WEBSOCKET DATA SOURCE -// ============================================================================ - -/** - * Data source for receiving real-time balance updates via WebSocket. - * - * This data source connects directly to BackendWebSocketService to receive - * push notifications for account balance changes. Unlike AccountsApiDataSource - * which polls for data, this provides instant updates. - * - * Uses Messenger pattern for all interactions: - * - Calls BackendWebSocketService methods via messenger actions - * - Exposes its own actions for AssetsController to call - * - Publishes events for AssetsController to subscribe to - * - * Actions exposed: - * - BackendWebsocketDataSource:getActiveChains - * - BackendWebsocketDataSource:subscribe - * - BackendWebsocketDataSource:unsubscribe - * - * Events published: - * - BackendWebsocketDataSource:activeChainsUpdated - * - BackendWebsocketDataSource:assetsUpdated - * - * Actions called (from BackendWebSocketService): - * - BackendWebSocketService:subscribe - * - BackendWebSocketService:getConnectionInfo - * - BackendWebSocketService:findSubscriptionsByChannelPrefix - * - BackendWebSocketService:addChannelCallback - * - BackendWebSocketService:removeChannelCallback - */ -const DEFAULT_CHAINS_REFRESH_INTERVAL_MS = 20 * 60 * 1000; // 20 minutes - -export class BackendWebsocketDataSource extends AbstractDataSource< - typeof CONTROLLER_NAME, - BackendWebsocketDataSourceState -> { - readonly #messenger: AssetsControllerMessenger; - - readonly #apiClient: ApiPlatformClient; - - readonly #onActiveChainsUpdated: ( - dataSourceName: string, - chains: ChainId[], - previousChains: ChainId[], - ) => void; - - readonly #getAssetType: ( - assetId: Caip19AssetId, - ) => 'native' | 'erc20' | 'spl'; - - /** Chains refresh timer */ - #chainsRefreshTimer: ReturnType | null = null; - - /** Chains the backend API reports as supported (preserved across disconnects). */ - #supportedChains: ChainId[] = []; - - /** Whether the WebSocket is currently connected. Chains are only claimed when true. */ - #isConnected = false; - - /** WebSocket subscriptions by our internal subscription ID */ - readonly #wsSubscriptions: Map = new Map(); - - /** Pending subscription requests to process when WebSocket connects */ - readonly #pendingSubscriptions: Map = new Map(); - - /** Store original subscription requests for reconnection */ - readonly #subscriptionRequests: Map = new Map(); - - /** Channels with registered BackendWebSocketService channel callbacks */ - readonly #registeredChannelCallbacks: Set = new Set(); - - /** Serializes subscribe/unsubscribe so account switches cannot interleave. */ - #subscribeLock: Promise = Promise.resolve(); - - constructor(options: BackendWebsocketDataSourceOptions) { - super(CONTROLLER_NAME, { - ...defaultState, - ...options.state, - }); - - this.#messenger = options.messenger; - this.#apiClient = options.queryApiClient; - this.#onActiveChainsUpdated = options.onActiveChainsUpdated; - this.#getAssetType = options.getAssetType; - - this.#subscribeToEvents(); - this.#initializeActiveChains().catch(console.error); - } - - // ============================================================================ - // INITIALIZATION - // ============================================================================ - - async #initializeActiveChains(): Promise { - try { - const chains = await this.#fetchActiveChains(); - this.#supportedChains = chains; - - // Only claim chains if the websocket is already connected. - // If not connected, chains stay unclaimed so AccountsApiDataSource - // can pick them up via polling. They'll be claimed on reconnect. - if (this.#isConnected) { - const previous = [...this.state.activeChains]; - this.updateActiveChains(chains, (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); - } - - this.#chainsRefreshTimer = setInterval(() => { - this.#refreshActiveChains().catch(console.error); - }, DEFAULT_CHAINS_REFRESH_INTERVAL_MS); - } catch (error) { - log('Failed to fetch active chains', error); - } - } - - async #refreshActiveChains(): Promise { - try { - const chains = await this.#fetchActiveChains(); - this.#supportedChains = chains; - - // Only update activeChains if connected; otherwise keep them unclaimed. - if (!this.#isConnected) { - return; - } - - const previousChains = new Set(this.state.activeChains); - const newChains = new Set(chains); - - const added = chains.filter((chain) => !previousChains.has(chain)); - const removed = Array.from(previousChains).filter( - (chain) => !newChains.has(chain), - ); - - if (added.length > 0 || removed.length > 0) { - const previous = [...this.state.activeChains]; - this.updateActiveChains(chains, (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); - } - } catch (error) { - log('Failed to refresh active chains', error); - } - } - - /** - * Re-fetch supported networks and refresh `activeChains` when connected. - * When disconnected, only `#supportedChains` is updated so reconnect can - * reclaim chains. Called on EVM network switch from AssetsController. - * - * @returns Resolves when supported networks have been re-fetched. - */ - refreshActiveChains(): Promise { - return this.#refreshActiveChains(); - } - - async #fetchActiveChains(): Promise { - const response = await this.#apiClient.accounts.fetchV2SupportedNetworks(); - return response.fullSupport.map(toChainId); - } - - #subscribeToEvents(): void { - type ConnectionStatePayload = { - state: WebSocketState; - [key: string]: unknown; - }; - // Listen for WebSocket connection state changes (event not in AssetsControllerEvents). - ( - this.#messenger as unknown as { - subscribe: (e: string, h: (p: ConnectionStatePayload) => void) => void; - } - ).subscribe( - 'BackendWebSocketService:connectionStateChanged', - (connectionInfo: ConnectionStatePayload) => { - if (connectionInfo.state === ('connected' as WebSocketState)) { - this.#isConnected = true; - this.#handleReconnect(); - } else if ( - connectionInfo.state === ('disconnected' as WebSocketState) - ) { - this.#isConnected = false; - this.#handleDisconnect(); - } - }, - ); - } - - /** - * Sync active chains from AccountsApiDataSource. - * When the data source invokes the onActiveChainsUpdated callback, the - * controller processes the active chains update (no messenger call; controller already updated). - * - * @param chains - Updated active chain IDs from AccountsApiDataSource. - */ - setActiveChainsFromAccountsApi(chains: ChainId[]): void { - this.updateActiveChains(chains, () => undefined); - } - - /** - * Handle WebSocket disconnection. - * Moves all active subscriptions to pending for re-subscription on reconnect. - */ - #handleDisconnect(): void { - log('WebSocket disconnected, releasing chains for fallback', { - activeSubscriptionCount: this.activeSubscriptions.size, - wsSubscriptionCount: this.#wsSubscriptions.size, - chainCount: this.state.activeChains.length, - }); - - // Move active subscriptions to pending for re-subscription - for (const [subscriptionId] of this.activeSubscriptions) { - const originalRequest = this.#subscriptionRequests.get(subscriptionId); - if (originalRequest) { - this.#pendingSubscriptions.set(subscriptionId, { - ...originalRequest, - isUpdate: false, - }); - } - } - - // Clear WebSocket subscriptions (server-side already cleared) - this.#wsSubscriptions.clear(); - - // Clear active subscriptions (they're no longer valid) - this.activeSubscriptions.clear(); - - // Release chains so the chain-claiming loop assigns them to - // AccountsApiDataSource (polling fallback) on the next #subscribeAssets. - const previous = [...this.state.activeChains]; - if (previous.length > 0) { - this.updateActiveChains([], (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); - } - } - - /** - * Handle WebSocket reconnection. - * Clears stale pending subscriptions and restores activeChains so the - * chain-claiming loop re-assigns them to this data source, triggering - * fresh subscriptions with current accounts and chains. - */ - #handleReconnect(): void { - log('WebSocket reconnected, reclaiming chains', { - supportedChainCount: this.#supportedChains.length, - pendingSubscriptionCount: this.#pendingSubscriptions.size, - }); - - // Discard stale pending subscriptions captured at disconnect time. - // The chain reclaim below triggers #onActiveChainsUpdated → - // #subscribeAssets() in AssetsController, which creates fresh - // subscriptions with current accounts and chains. Processing the - // stale pending entries afterwards would overwrite those with - // outdated request data. - this.#pendingSubscriptions.clear(); - - if (this.#supportedChains.length > 0) { - const previous = [...this.state.activeChains]; - this.updateActiveChains(this.#supportedChains, (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); - } - } - - // ============================================================================ - // ACTIVE CHAINS - // ============================================================================ - - /** - * Update active chains when AccountsApiDataSource reports new supported chains. - * - * @param chains - Array of supported chain IDs. - */ - updateSupportedChains(chains: ChainId[]): void { - const previous = [...this.state.activeChains]; - this.updateActiveChains(chains, (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); - } - - // ============================================================================ - // SUBSCRIBE - // ============================================================================ - - async subscribe(subscriptionRequest: SubscriptionRequest): Promise { - const previousLock = this.#subscribeLock; - let releaseLock: () => void = () => undefined; - this.#subscribeLock = new Promise((resolve) => { - releaseLock = resolve; - }); - - await previousLock; - try { - await this.#subscribeInternal(subscriptionRequest); - } finally { - releaseLock(); - } - } - - async #subscribeInternal( - subscriptionRequest: SubscriptionRequest, - ): Promise { - const { request, subscriptionId, isUpdate } = subscriptionRequest; - - // Filter to active chains only - const chainsToSubscribe = request.chainIds.filter((chainId) => - this.state.activeChains.includes(chainId), - ); - - const addresses = request.accountsWithSupportedChains.map( - (a) => a.account.address, - ); - - if (addresses.length === 0) { - return; - } - - // Check WebSocket connection status - try { - const connectionInfo = this.#messenger.call( - 'BackendWebSocketService:getConnectionInfo', - ); - if (connectionInfo.state !== ('connected' as WebSocketState)) { - // Store the subscription request to process when WebSocket connects - this.#pendingSubscriptions.set(subscriptionId, subscriptionRequest); - return; - } - } catch { - // Store anyway - will be processed when we can connect - this.#pendingSubscriptions.set(subscriptionId, subscriptionRequest); - return; - } - - // Remove from pending if it was there (we're processing it now) - this.#pendingSubscriptions.delete(subscriptionId); - - // Handle subscription update - if (isUpdate) { - const existing = this.activeSubscriptions.get(subscriptionId); - if (existing) { - // Check if accounts changed - if so, we need to re-subscribe to different channels - const existingAddresses = existing.addresses ?? []; - const addressesChanged = haveAddressesChanged( - addresses, - existingAddresses, - ); - - if (!addressesChanged) { - // Only chains changed - update chains, request, and callback - existing.chains = chainsToSubscribe; - existing.onAssetsUpdate = subscriptionRequest.onAssetsUpdate; - this.#subscriptionRequests.set(subscriptionId, subscriptionRequest); - return; - } - // Accounts changed - fall through to re-subscribe with new channels - } - } - - // Clean up existing subscription if any (inline teardown — subscribe holds the lock) - await this.#teardownSubscription(subscriptionId); - - // Always subscribe to eip155 and solana account activity, plus any namespaces from requested chains - const namespaces = getNamespacesForAccountActivity(chainsToSubscribe); - - // Build channel names: use namespace-appropriate address per account (eip155 = hex, solana = base58) - const channels: string[] = []; - for (const namespace of namespaces) { - for (const { account } of request.accountsWithSupportedChains) { - const address = getAddressForAccountActivity(account, namespace); - if (address) { - channels.push(buildAccountActivityChannel(namespace, address)); - } - } - } - - if (channels.length === 0) { - return; - } - - try { - // Register request/callback before awaiting server subscribe so notifications - // that arrive during the subscribe handshake are not dropped. - this.#subscriptionRequests.set(subscriptionId, subscriptionRequest); - this.activeSubscriptions.set(subscriptionId, { - cleanup: () => { - this.#teardownSubscription(subscriptionId).catch(() => undefined); - }, - chains: chainsToSubscribe, - addresses, - onAssetsUpdate: subscriptionRequest.onAssetsUpdate, - }); - - // Create WebSocket subscription - const wsSubscription = await this.#messenger.call( - 'BackendWebSocketService:subscribe', - { - channels, - channelType: CHANNEL_TYPE, - callback: (notification: ServerNotificationMessage) => { - this.#handleNotification(notification, subscriptionId); - }, - }, - ); - - this.#wsSubscriptions.set(subscriptionId, wsSubscription); - - try { - this.#registerChannelCallbacks(subscriptionId, channels); - } catch (channelCallbackError) { - log( - 'Channel callback registration failed; ws subscription still active', - { subscriptionId, error: channelCallbackError }, - ); - } - } catch (error) { - this.activeSubscriptions.delete(subscriptionId); - this.#subscriptionRequests.delete(subscriptionId); - log('WebSocket subscription FAILED', { - subscriptionId, - error, - chains: chainsToSubscribe, - }); - } - } - - // ============================================================================ - // NOTIFICATION HANDLING - // ============================================================================ - - #handleNotification( - notification: ServerNotificationMessage, - subscriptionId: string, - ): void { - try { - const activityMessage = - notification.data as unknown as AccountActivityMessage; - - const storedSubscription = this.#subscriptionRequests.get(subscriptionId); - const request = storedSubscription?.request; - const onAssetsUpdate = - this.activeSubscriptions.get(subscriptionId)?.onAssetsUpdate ?? - storedSubscription?.onAssetsUpdate; - - if (!request) { - return; - } - - const { address, tx, updates } = activityMessage; - - if (!address || !tx || !updates) { - return; - } - - // Extract chain ID from transaction (CAIP-2 format, e.g., "eip155:8453") - const chainId = tx.chain as ChainId; - - // Find matching account in request (eip155: case-insensitive hex; solana: exact base58) - const account = request.accountsWithSupportedChains - .map((entry) => entry.account) - .find((a) => - a.address.startsWith('0x') - ? a.address.toLowerCase() === address.toLowerCase() - : a.address === address, - ); - if (!account) { - return; - } - const accountId = account.id; - - // Process all balance updates from the activity message - const response = this.#processBalanceUpdates(updates, chainId, accountId); - - const balanceEntries = response.assetsBalance?.[accountId] ?? {}; - const hasBalances = Object.keys(balanceEntries).length > 0; - - if (hasBalances && onAssetsUpdate) { - Promise.resolve(onAssetsUpdate(response, request)).catch((error) => { - console.error(error); - }); - } - } catch (error) { - log('Error handling notification', error); - } - } - - /** - * Process balance updates from AccountActivityMessage. - * Each update contains asset info, post-transaction balance, and transfer details. - * - * @param updates - Array of balance updates from the activity message. - * @param _chainId - The chain ID (unused but kept for context). - * @param accountId - The account ID to process updates for. - * @returns DataResponse containing processed balance and metadata. - */ - #processBalanceUpdates( - updates: BalanceUpdate[], - _chainId: ChainId, - accountId: string, - ): DataResponse { - return processAccountActivityBalanceUpdates(updates, accountId, (assetId) => - this.#getAssetType(assetId), - ); - } - - // ============================================================================ - // UNSUBSCRIBE - // ============================================================================ - - /** - * Unsubscribe and await server-side teardown so a re-subscribe does not race - * with stale subscription IDs on incoming notifications. - * - * @param subscriptionId - The ID of the subscription to cancel. - */ - async unsubscribe(subscriptionId: string): Promise { - const previousLock = this.#subscribeLock; - let releaseLock: () => void = () => undefined; - this.#subscribeLock = new Promise((resolve) => { - releaseLock = resolve; - }); - - await previousLock; - try { - await this.#teardownSubscription(subscriptionId); - } finally { - releaseLock(); - } - } - - async #teardownSubscription(subscriptionId: string): Promise { - const wsSub = this.#wsSubscriptions.get(subscriptionId); - - if (wsSub) { - const channels = [...wsSub.channels]; - try { - await wsSub.unsubscribe(); - } catch (unsubErr: unknown) { - log('Error unsubscribing', { subscriptionId, error: unsubErr }); - } - this.#wsSubscriptions.delete(subscriptionId); - this.#removeChannelCallbacks(channels); - } - - this.#subscriptionRequests.delete(subscriptionId); - this.activeSubscriptions.delete(subscriptionId); - } - - #registerChannelCallbacks(subscriptionId: string, channels: string[]): void { - for (const channel of channels) { - this.#unregisterChannelCallback(channel); - - try { - this.#messenger.call('BackendWebSocketService:addChannelCallback', { - channelName: channel, - callback: (notification: ServerNotificationMessage) => { - this.#handleNotification(notification, subscriptionId); - }, - }); - this.#registeredChannelCallbacks.add(channel); - } catch { - // Channel callbacks are optional; ws subscription still works without them. - } - } - } - - #unregisterChannelCallback(channel: string): void { - if (!this.#registeredChannelCallbacks.has(channel)) { - return; - } - - try { - this.#messenger.call( - 'BackendWebSocketService:removeChannelCallback', - channel, - ); - } catch { - // Best-effort cleanup when the channel callback was never registered. - } - - this.#registeredChannelCallbacks.delete(channel); - } - - #removeChannelCallbacks(channels: string[]): void { - for (const channel of channels) { - this.#unregisterChannelCallback(channel); - } - } - - // ============================================================================ - // CLEANUP - // ============================================================================ - - destroy(): void { - if (this.#chainsRefreshTimer) { - clearInterval(this.#chainsRefreshTimer); - this.#chainsRefreshTimer = null; - } - - const subscriptionIds = [ - ...new Set([ - ...this.#wsSubscriptions.keys(), - ...this.activeSubscriptions.keys(), - ]), - ]; - for (const subscriptionId of subscriptionIds) { - this.#teardownSubscription(subscriptionId).catch(() => undefined); - } - - // Clean up base class subscriptions (no-op if already torn down) - super.destroy(); - } -} - -// ============================================================================ -// FACTORY FUNCTION -// ============================================================================ - -/** - * Creates a BackendWebsocketDataSource instance. - * - * @param options - Configuration options for the data source. - * @returns A new BackendWebsocketDataSource instance. - */ -export function createBackendWebsocketDataSource( - options: BackendWebsocketDataSourceOptions, -): BackendWebsocketDataSource { - return new BackendWebsocketDataSource(options); -} diff --git a/packages/assets-controller/src/data-sources/index.ts b/packages/assets-controller/src/data-sources/index.ts index be96ba70d72..04758c57281 100644 --- a/packages/assets-controller/src/data-sources/index.ts +++ b/packages/assets-controller/src/data-sources/index.ts @@ -13,13 +13,12 @@ export { } from './AccountsApiDataSource.js'; export { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, - type BackendWebsocketDataSourceOptions, - type BackendWebsocketDataSourceState, - type BackendWebsocketDataSourceAllowedActions, - type BackendWebsocketDataSourceAllowedEvents, -} from './BackendWebsocketDataSource.js'; + AccountActivityDataSource, + createAccountActivityDataSource, + type AccountActivityDataSourceOptions, + type AccountActivityDataSourceState, + type AccountActivityDataSourceAllowedEvents, +} from './AccountActivityDataSource.js'; export { RpcDataSource, diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 6f9185b7bda..8a59fe53714 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -126,17 +126,6 @@ export type { AccountsApiDataSourceState, } from './data-sources/index.js'; -// Data sources - BackendWebsocket -export { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, -} from './data-sources/index.js'; - -export type { - BackendWebsocketDataSourceOptions, - BackendWebsocketDataSourceState, -} from './data-sources/index.js'; - // Data sources - RPC export { RpcDataSource, createRpcDataSource } from './data-sources/index.js'; diff --git a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts index 2fdcb05d5f6..0b3613fe2d0 100644 --- a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts +++ b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts @@ -178,7 +178,7 @@ describe('CustomAssetGraduationMiddleware', () => { }); it('does not graduate when the decimal amount is zero', async () => { - // Both AccountsApi and BackendWebsocketDataSource emit human-readable + // Both AccountsApi and AccountActivityDataSource emit human-readable // decimal strings, so "0.0" must be treated the same as "0". const { middleware, context, removeCustomAsset } = setup({ [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], @@ -424,7 +424,7 @@ describe('CustomAssetGraduationMiddleware', () => { }); it('graduates a custom asset when the response uses a non-checksummed (lowercase) address', async () => { - // Regression: BackendWebsocketDataSource does not normalize asset IDs, + // Regression: AccountActivityDataSource does not normalize asset IDs, // so balances may arrive with lowercase addresses while customAssets // state stores the checksummed form. Graduation must be robust to that. const checksummedCustomAsset = diff --git a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts index 9319a2f6d87..17070e3cdde 100644 --- a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts +++ b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts @@ -83,7 +83,7 @@ export class CustomAssetGraduationMiddleware { // customAssets state is stored with checksummed/normalized asset IDs. // AccountsApiDataSource normalizes its response IDs, but - // BackendWebsocketDataSource does not — so we normalize the response + // AccountActivityDataSource does not — so we normalize the response // side here to make the comparison robust to lower-case addresses // delivered over the websocket. const customSet = new Set(customForAccount); diff --git a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.test.ts b/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.test.ts deleted file mode 100644 index 01f1d82b47a..00000000000 --- a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.test.ts +++ /dev/null @@ -1,39 +0,0 @@ -import type { BalanceUpdate } from '@metamask/core-backend'; - -import type { Caip19AssetId } from '../types.js'; -import { processAccountActivityBalanceUpdates } from './processAccountActivityBalanceUpdates.js'; - -describe('processAccountActivityBalanceUpdates', () => { - it('converts hex postBalance to human-readable amount', () => { - const accountId = 'account-1'; - const assetId = 'eip155:42161/slip44:60' as Caip19AssetId; - const updates = [ - { - asset: { - fungible: true, - type: assetId, - unit: 'ETH', - decimals: 18, - }, - postBalance: { amount: '0x10aa6d94e80' }, - transfers: [], - }, - ] as BalanceUpdate[]; - - const response = processAccountActivityBalanceUpdates( - updates, - accountId, - () => 'native', - ); - - expect(response.updateMode).toBe('merge'); - expect(response.assetsBalance?.[accountId]?.[assetId]).toStrictEqual({ - amount: '0.00000114526056', - }); - expect(response.assetsInfo?.[assetId]).toMatchObject({ - type: 'native', - symbol: 'ETH', - decimals: 18, - }); - }); -}); diff --git a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts b/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts deleted file mode 100644 index b6367ff9b15..00000000000 --- a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts +++ /dev/null @@ -1,78 +0,0 @@ -import type { BalanceUpdate } from '@metamask/core-backend'; -import BigNumberJS from 'bignumber.js'; - -import type { - AssetBalance, - AssetMetadata, - Caip19AssetId, - DataResponse, -} from '../types.js'; - -/** - * Convert AccountActivityMessage balance updates into a {@link DataResponse} - * for AssetsController (same shape as BackendWebsocketDataSource). - * - * @param updates - Balance updates from account-activity websocket payload. - * @param accountId - Internal account UUID. - * @param getAssetType - Resolver for asset metadata type. - * @returns DataResponse with merge mode when balances are present. - */ -export function processAccountActivityBalanceUpdates( - updates: BalanceUpdate[], - accountId: string, - getAssetType: (assetId: Caip19AssetId) => 'native' | 'erc20' | 'spl', -): DataResponse { - const assetsBalance = Object.create(null) as Record< - string, - Record - >; - assetsBalance[accountId] = Object.create(null) as Record< - Caip19AssetId, - AssetBalance - >; - const assetsMetadata = Object.create(null) as Record< - Caip19AssetId, - AssetMetadata - >; - - for (const update of updates) { - const { asset, postBalance } = update; - - if (!asset || !postBalance) { - continue; - } - - const assetId = asset.type as Caip19AssetId; - - if (asset.decimals === undefined) { - continue; - } - - const rawBalanceStr = postBalance.amount.startsWith('0x') - ? BigInt(postBalance.amount).toString() - : postBalance.amount; - - const humanReadableAmount = new BigNumberJS(rawBalanceStr) - .dividedBy(new BigNumberJS(10).pow(asset.decimals)) - .toFixed(); - - assetsBalance[accountId][assetId] = { - amount: humanReadableAmount, - }; - - assetsMetadata[assetId] = { - type: getAssetType(assetId), - symbol: asset.unit, - name: asset.unit, - decimals: asset.decimals, - }; - } - - const response: DataResponse = { updateMode: 'merge' }; - if (Object.keys(assetsBalance[accountId]).length > 0) { - response.assetsBalance = assetsBalance; - response.assetsInfo = assetsMetadata; - } - - return response; -}