From cfc45457d879391eda55a2bba0c5ad0b1dcb1023 Mon Sep 17 00:00:00 2001 From: Kriys94 Date: Wed, 15 Jul 2026 12:52:05 +0200 Subject: [PATCH 1/2] fix(assets-controller): consume account activity --- .../src/AssetsController.test.ts | 110 +++++- .../assets-controller/src/AssetsController.ts | 197 ++++++---- .../data-sources/AccountActivityDataSource.ts | 342 ++++++++++++++++++ .../AccountsApiDataSource.test.ts | 129 +++++++ .../src/data-sources/AccountsApiDataSource.ts | 42 ++- .../src/data-sources/RpcDataSource.test.ts | 75 +++- .../src/data-sources/RpcDataSource.ts | 56 ++- .../src/data-sources/index.ts | 8 + packages/assets-controller/src/utils/index.ts | 1 + .../src/utils/subscriptionScope.test.ts | 124 +++++++ .../src/utils/subscriptionScope.ts | 45 +++ 11 files changed, 1042 insertions(+), 87 deletions(-) create mode 100644 packages/assets-controller/src/data-sources/AccountActivityDataSource.ts create mode 100644 packages/assets-controller/src/utils/subscriptionScope.test.ts create mode 100644 packages/assets-controller/src/utils/subscriptionScope.ts diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 6916fd97a9b..f50b1e98f27 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -86,6 +86,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; @@ -1606,14 +1612,81 @@ describe('AssetsController', () => { }); describe('handleActiveChainsUpdate', () => { - it('re-subscribes assets when chains are added', async () => { + it('re-subscribes assets when chains are added, debounced and jittered', async () => { await withController(async ({ controller }) => { - const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + 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(); - expect(subscribeSpy).toHaveBeenCalledTimes(1); + // 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('re-subscribes promptly (no jitter) when a chain removal pre-empts a pending jittered addition', async () => { + await withController(async ({ controller }) => { + jest.useFakeTimers(); + const randomSpy = jest.spyOn(Math, 'random').mockReturnValue(0.9); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + // Addition schedules a jittered re-subscribe... + onActiveChainsUpdated('TestDataSource', ['eip155:1'], []); + // ...then a removal arrives; it must pre-empt to the prompt window. + onActiveChainsUpdated('TestDataSource', [], ['eip155:137']); + + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_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 }) => { + jest.useFakeTimers(); + 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'], + ); + + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_MS); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } }); }); @@ -1642,12 +1715,18 @@ describe('AssetsController', () => { it('re-subscribes assets when chains are removed', async () => { await withController(async ({ controller }) => { - const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); + jest.useFakeTimers(); + try { + const subscribeSpy = jest.spyOn(controller, 'subscribeAssetsPrice'); - const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); - onActiveChainsUpdated('TestDataSource', [], ['eip155:1']); + const onActiveChainsUpdated = controller.getOnActiveChainsUpdated(); + onActiveChainsUpdated('TestDataSource', [], ['eip155:1']); - expect(subscribeSpy).toHaveBeenCalledTimes(1); + jest.advanceTimersByTime(RESUBSCRIBE_DEBOUNCE_MS); + expect(subscribeSpy).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } }); }); @@ -2233,8 +2312,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 0ed051dcce9..82689ee7a80 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -16,9 +16,9 @@ import type { TraceCallback } from '@metamask/controller-utils'; import type { ApiPlatformClient, AccountActivityServiceBalanceUpdatedEvent, + AccountActivityServiceStatusChangedEvent, BackendWebSocketServiceActions, BackendWebSocketServiceEvents, - BalanceUpdate, SupportedCurrency, } from '@metamask/core-backend'; import type { @@ -78,6 +78,7 @@ import type { } from './data-sources/AbstractDataSource'; import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource'; import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource'; +import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource'; import { BackendWebsocketDataSource } from './data-sources/BackendWebsocketDataSource'; import { shouldSkipNativeForCaipChainId } from './data-sources/evm-rpc-services/utils/assets'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource'; @@ -147,7 +148,6 @@ import type { } from './utils'; import { ZERO_ADDRESS } from './utils/constants'; import { pickRpcCustomAssetsSupplement } from './utils/customAssetsRpcSupplement'; -import { processAccountActivityBalanceUpdates } from './utils/processAccountActivityBalanceUpdates'; const NATIVE_ASSETS_QUERY_KEY = ['nativeAssets']; @@ -197,6 +197,26 @@ 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 when a data source + * *claims* new chains (e.g. WebSocket `statusChanged: 'up'`). Staggers the + * resulting WS-subscribe fan-out across clients that all receive the same + * backend broadcast, avoiding a thundering herd. Only applied to chain + * *additions*: additions are non-urgent (polling still covers the chain until + * re-subscribe), whereas removals are re-subscribed promptly so polling fallback + * resumes without a data gap. + */ +const RESUBSCRIBE_JITTER_MS = 5000; + // ============================================================================ // TRACE NAMES — used in Sentry spans (search these strings in Discover) // ============================================================================ @@ -352,8 +372,9 @@ type AllowedEvents = | SnapControllerSnapInstalledEvent // BackendWebsocketDataSource | BackendWebSocketServiceEvents - // AccountActivityService (real-time balance updates for unified assets) - | AccountActivityServiceBalanceUpdatedEvent; + // AccountActivityService (real-time balance updates + chain status for unified assets) + | AccountActivityServiceBalanceUpdatedEvent + | AccountActivityServiceStatusChangedEvent; export type AssetsControllerMessenger = Messenger< typeof CONTROLLER_NAME, @@ -745,6 +766,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(); @@ -780,6 +815,8 @@ export class AssetsController extends BaseController< readonly #backendWebsocketDataSource: BackendWebsocketDataSource; + readonly #accountActivityDataSource: AccountActivityDataSource; + readonly #accountsApiDataSource: AccountsApiDataSource; readonly #snapDataSource: SnapDataSource; @@ -900,6 +937,12 @@ export class AssetsController extends BaseController< getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => this.#getAssetType(assetId), }); + this.#accountActivityDataSource = new AccountActivityDataSource({ + messenger: this.messenger, + onActiveChainsUpdated: this.#onActiveChainsUpdated, + getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => + this.#getAssetType(assetId), + }); this.#accountsApiDataSource = new AccountsApiDataSource({ messenger: this.messenger, queryApiClient, @@ -1177,15 +1220,8 @@ export class AssetsController extends BaseController< }, ); - // 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); - }, - ); + // Real-time post-tx balances and chain status from AccountActivityService are + // consumed by AccountActivityDataSource (subscribed in #subscribeAssets). } #onUnapprovedTransactionAdded(transactionMeta: TransactionMeta): void { @@ -1244,49 +1280,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. @@ -1447,12 +1440,23 @@ export class AssetsController extends BaseController< const addedChains = activeChains.filter((ch) => !previousSet.has(ch)); const removedChains = previous.filter((ch) => !activeChains.includes(ch)); - 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(); + // 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. + if (removedChains.length > 0) { + // Urgent: a data source dropped chains (e.g. WebSocket went down). Until + // re-subscribe hands those chains to a polling source there is a data gap, + // so run promptly without jitter. + this.#scheduleSubscribeAssets(); + } else if (addedChains.length > 0) { + // Non-urgent: a data source claimed new chains (polling still covers them + // until re-subscribe). Add jitter to stagger the resulting WS-subscribe + // fan-out across clients receiving the same backend broadcast. + this.#scheduleSubscribeAssets({ jitter: true }); } } @@ -2943,6 +2947,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(); @@ -2952,6 +2963,7 @@ export class AssetsController extends BaseController< // every source that may have been subscribed. const allSources = [ ...this.#allBalanceDataSources, + this.#accountActivityDataSource, this.#stakedBalanceDataSource, ]; const subscriptionKeys = [...this.#activeSubscriptions.keys()]; @@ -2988,6 +3000,53 @@ 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. + * + * When `jitter` is set, a random delay up to {@link RESUBSCRIBE_JITTER_MS} is + * added to stagger the WS-subscribe fan-out across clients (used for chain + * additions). A later non-jittered call pre-empts a pending jittered one so an + * urgent chain removal is not held back by an addition's jitter window. + * + * @param options - Scheduling options. + * @param options.jitter - Add random jitter to the delay (for chain additions). + */ + #scheduleSubscribeAssets({ jitter = false }: { jitter?: boolean } = {}): void { + const delay = jitter + ? RESUBSCRIBE_DEBOUNCE_MS + Math.floor(Math.random() * RESUBSCRIBE_JITTER_MS) + : RESUBSCRIBE_DEBOUNCE_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; + try { + this.#subscribeAssets(); + } catch (error) { + log('Failed to run coalesced re-subscribe', error); + } + }, delay); + } + /** * Subscribe to asset updates for all selected accounts. */ @@ -3084,6 +3143,17 @@ export class AssetsController extends BaseController< chainToAccounts, rpcAssignedChains, ); + + // AccountActivityDataSource consumes real-time AccountActivityService events. + // It does not participate in chain-claiming; it is subscribed with all + // selected accounts and enabled chains so it can route `balanceUpdated` + // events to the matching account. Chain up/down is managed internally from + // `statusChanged` (see AccountActivityDataSource). + this.#subscribeDataSource( + this.#accountActivityDataSource, + accounts, + chainIds, + ); } /** @@ -3682,7 +3752,7 @@ export class AssetsController extends BaseController< const shouldGraduateCustomAssets = sourceId === 'AccountsApiDataSource' || sourceId === 'BackendWebsocketDataSource' || - sourceId === 'AccountActivityService'; + sourceId === 'AccountActivityDataSource'; const enrichmentSources: AssetsDataSource[] = [ ...(shouldGraduateCustomAssets @@ -3734,6 +3804,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/data-sources/AccountActivityDataSource.ts b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts new file mode 100644 index 00000000000..6bb227b48f7 --- /dev/null +++ b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts @@ -0,0 +1,342 @@ +import type { + AccountActivityServiceBalanceUpdatedEvent, + AccountActivityServiceStatusChangedEvent, + BalanceUpdate, +} from '@metamask/core-backend'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; +import { isCaipChainId } from '@metamask/utils'; + +import type { AssetsControllerMessenger } from '../AssetsController'; +import { projectLogger, createModuleLogger } from '../logger'; +import type { ChainId, Caip19AssetId, DataRequest } from '../types'; +import { processAccountActivityBalanceUpdates } from '../utils/processAccountActivityBalanceUpdates'; +import { AbstractDataSource } from './AbstractDataSource'; +import type { + DataSourceState, + SubscriptionRequest, +} from './AbstractDataSource'; + +// ============================================================================ +// CONSTANTS +// ============================================================================ + +const CONTROLLER_NAME = 'AccountActivityDataSource'; + +const log = createModuleLogger(projectLogger, CONTROLLER_NAME); + +// ============================================================================ +// 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'; + state?: Partial; +}; + +// ============================================================================ +// ACCOUNT ACTIVITY DATA SOURCE +// ============================================================================ + +/** + * Data source that consumes real-time updates from `AccountActivityService`. + * + * Unlike {@link BackendWebsocketDataSource}, which owns its own WebSocket + * channel subscriptions, this data source is a thin consumer of the two + * high-level events that `AccountActivityService` publishes: + * + * - `AccountActivityService:balanceUpdated` — post-transaction balances for the + * subscribed account(s). The address is resolved against the accounts in the + * active subscription(s), transformed into a {@link DataResponse} (via + * {@link processAccountActivityBalanceUpdates}), and pushed to the controller + * through the subscription's `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` also + * are never seeded from the accounts API the way `BackendWebsocketDataSource` + * does; 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) and is wired into the controller's balance + * subscription flow so incoming balance updates can be routed to the right + * account. + */ +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'; + + /** Stored subscription requests (used to route incoming balance updates). */ + readonly #subscriptionRequests: Map = new Map(); + + /** 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.#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 + // ============================================================================ + + async subscribe(subscriptionRequest: SubscriptionRequest): Promise { + const { subscriptionId, request } = subscriptionRequest; + + // Store the request so incoming `balanceUpdated` events can be routed to the + // matching account and reported through its `onAssetsUpdate` callback. + this.#subscriptionRequests.set(subscriptionId, subscriptionRequest); + this.activeSubscriptions.set(subscriptionId, { + cleanup: () => { + this.#subscriptionRequests.delete(subscriptionId); + }, + chains: request.chainIds, + addresses: request.accountsWithSupportedChains.map( + (entry) => entry.account.address, + ), + onAssetsUpdate: subscriptionRequest.onAssetsUpdate, + }); + } + + // ============================================================================ + // BALANCE UPDATES + // ============================================================================ + + #onBalanceUpdated({ + address, + chain, + updates, + }: { + address: string; + chain: string; + updates: BalanceUpdate[]; + }): void { + try { + if (!address || !chain || !updates || updates.length === 0) { + return; + } + + const match = this.#findAccountForAddress(address); + if (!match) { + return; + } + + const { account, onAssetsUpdate } = match; + 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(onAssetsUpdate(response, request)).catch((error) => { + log('Failed to report balance update', { error }); + }); + } catch (error) { + log('Error handling balance update', error); + } + } + + /** + * Find the internal account matching the given activity address across all + * active subscription requests, along with the callback used to report its + * updates. EVM addresses are matched case-insensitively; other namespaces are + * matched exactly. + * + * @param address - The account address from the activity message. + * @returns The matching account and its `onAssetsUpdate` callback, or null. + */ + #findAccountForAddress(address: string): { + account: InternalAccount; + onAssetsUpdate: SubscriptionRequest['onAssetsUpdate']; + } | null { + const isEvm = address.startsWith('0x'); + + for (const subscriptionRequest of this.#subscriptionRequests.values()) { + const account = subscriptionRequest.request.accountsWithSupportedChains + .map((entry) => entry.account) + .find((candidate) => + isEvm + ? candidate.address.toLowerCase() === address.toLowerCase() + : candidate.address === address, + ); + + if (account) { + return { + account, + onAssetsUpdate: subscriptionRequest.onAssetsUpdate, + }; + } + } + + return null; + } + + // ============================================================================ + // 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`. + */ + 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; + + this.#subscriptionRequests.clear(); + + 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/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index f01829bf159..2ac0cf1f3ce 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -774,6 +774,135 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); + it('skips the immediate fetch when re-subscribing with an unchanged scope within one poll interval', async () => { + const { controller, apiClient, assetsUpdateHandler } = + await setupController(); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); + + // Re-subscribe with the identical scope (e.g. a chain flapped back from the + // WebSocket source). The immediate fetch is skipped as redundant. + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); + + controller.destroy(); + }); + + it('performs the immediate fetch when re-subscribing with forceUpdate even if the scope is unchanged', async () => { + const { controller, apiClient, assetsUpdateHandler } = + await setupController(); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest({ forceUpdate: true }), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(2); + + controller.destroy(); + }); + + it('performs the immediate fetch when re-subscribing with a different scope', async () => { + const { controller, apiClient, assetsUpdateHandler } = + await setupController(); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest({ chainIds: [CHAIN_MAINNET] }), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON] }), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(2); + + controller.destroy(); + }); + + it('performs the immediate fetch when re-subscribing after the previous fetch has gone stale', async () => { + const { controller, apiClient, assetsUpdateHandler } = + await setupController(); + + const nowSpy = jest.spyOn(Date, 'now'); + const startTime = 1_000_000; + nowSpy.mockReturnValue(startTime); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + + // Advance the clock past the default poll interval (30s) so the previous + // fetch is no longer fresh. + nowSpy.mockReturnValue(startTime + 30_001); + + await controller.subscribe({ + subscriptionId: 'sub-1', + request: createDataRequest(), + isUpdate: false, + onAssetsUpdate: assetsUpdateHandler, + }); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(2); + + nowSpy.mockRestore(); + controller.destroy(); + }); + it('subscribe does nothing when no chains', async () => { const { controller, assetsUpdateHandler } = await setupController(); diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index 6ac6167b0e4..73f4cdb3e59 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -21,7 +21,11 @@ import type { Middleware, AssetsControllerStateInternal, } from '../types'; -import { fetchWithTimeout, normalizeAssetId } from '../utils'; +import { + computeSubscriptionScopeKey, + fetchWithTimeout, + normalizeAssetId, +} from '../utils'; import type { DataSourceState, SubscriptionRequest, @@ -216,6 +220,15 @@ export class AccountsApiDataSource extends AbstractDataSource< /** State accessor from subscriptions (for filtering when tokenDetectionEnabled is false) */ #getAssetsState?: () => AssetsControllerStateInternal; + /** + * Scope key + timestamp of the last fetch, preserved across teardown. Used to + * skip the redundant immediate fetch when the subscription is re-created for + * the same scope while its last fetch is still fresh (e.g. a chain flapping + * between the WebSocket and this source). This source only ever has a single + * subscription, so a single record suffices. + */ + #lastFetch: { scopeKey: string; at: number } | null = null; + constructor(options: AccountsApiDataSourceOptions) { super(CONTROLLER_NAME, { ...defaultState, @@ -752,6 +765,11 @@ export class AccountsApiDataSource extends AbstractDataSource< await this.unsubscribe(subscriptionId); const pollInterval = request.updateInterval ?? this.#pollInterval; + const scopeKey = computeSubscriptionScopeKey( + request, + chainsToSubscribe, + pollInterval, + ); // Create poll function for this subscription const pollFn = async (): Promise => { @@ -761,6 +779,8 @@ export class AccountsApiDataSource extends AbstractDataSource< return; } + this.#lastFetch = { scopeKey, at: Date.now() }; + // Use stored request (which gets updated on account changes) const fetchResponse = await this.fetch({ ...subscription.request, @@ -789,6 +809,24 @@ export class AccountsApiDataSource extends AbstractDataSource< onAssetsUpdate: subscriptionRequest.onAssetsUpdate, }); + // Skip the immediate fetch when this subscription was just re-created for + // the same scope and its last fetch is still fresh (within one poll + // interval). This avoids redundant identical requests when a chain flaps + // between the WebSocket source and this polling source. `forceUpdate` + // always fetches (it is excluded from the scope key). The scheduled poll + // above still refreshes on the next interval tick. + const isRedundantImmediateFetch = + !request.forceUpdate && + this.#lastFetch?.scopeKey === scopeKey && + Date.now() - this.#lastFetch.at < pollInterval; + + if (isRedundantImmediateFetch) { + log('Skipping redundant immediate fetch on re-subscribe', { + subscriptionId, + }); + return; + } + // Initial fetch await pollFn(); } @@ -803,6 +841,8 @@ export class AccountsApiDataSource extends AbstractDataSource< clearInterval(this.#chainsRefreshTimer); } + this.#lastFetch = null; + // Clean up subscriptions super.destroy(); } diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts index c2ff9044d13..2ae3e8f6b56 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts @@ -1276,11 +1276,49 @@ describe('RpcDataSource', () => { }); }); - it('updates existing subscription when isUpdate true', async () => { + it('restarts polling on update when the scope changes', async () => { const balanceStartSpy = jest.spyOn( BalanceFetcher.prototype, 'startPolling', ); + const secondAccount = createMockInternalAccount({ + id: 'account-2', + address: '0x9999999999999999999999999999999999999999', + }); + await withController(async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + balanceStartSpy.mockClear(); + + // Adding a second account changes the subscription scope, so polling is + // restarted for the new set of accounts. + await controller.subscribe({ + request: createDataRequest({ + accounts: [createMockInternalAccount(), secondAccount], + }), + subscriptionId: 'test-sub', + isUpdate: true, + onAssetsUpdate: jest.fn(), + }); + expect(balanceStartSpy).toHaveBeenCalledWith( + expect.objectContaining({ accountId: 'account-2' }), + ); + }); + }); + + it('skips restarting polling when re-subscribing with an unchanged scope within one interval', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const balanceStopSpy = jest.spyOn( + BalanceFetcher.prototype, + 'stopPollingByPollingToken', + ); await withController(async ({ controller }) => { await controller.subscribe({ request: createDataRequest(), @@ -1288,13 +1326,46 @@ describe('RpcDataSource', () => { isUpdate: false, onAssetsUpdate: jest.fn(), }); + balanceStartSpy.mockClear(); + balanceStopSpy.mockClear(); + + // Re-subscribing with the identical scope (e.g. a chain flapping back + // from the WebSocket source) is a no-op: polling is neither stopped nor + // restarted, so no redundant immediate RPC fetch is triggered. + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: true, + onAssetsUpdate: jest.fn(), + }); + expect(balanceStartSpy).not.toHaveBeenCalled(); + expect(balanceStopSpy).not.toHaveBeenCalled(); + }); + }); + + it('restarts polling when re-subscribing with forceUpdate even if the scope is unchanged', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + await withController(async ({ controller }) => { await controller.subscribe({ request: createDataRequest(), subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + balanceStartSpy.mockClear(); + + await controller.subscribe({ + request: createDataRequest({ forceUpdate: true }), + subscriptionId: 'test-sub', isUpdate: true, onAssetsUpdate: jest.fn(), }); - expect(balanceStartSpy).toHaveBeenCalledTimes(2); + expect(balanceStartSpy).toHaveBeenCalledWith( + expect.objectContaining({ accountId: MOCK_ACCOUNT_ID }), + ); }); }); diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.ts b/packages/assets-controller/src/data-sources/RpcDataSource.ts index d28af6e3fda..cd63f841c21 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.ts @@ -35,7 +35,7 @@ import type { DataResponse, Middleware, } from '../types'; -import { normalizeAssetId } from '../utils'; +import { computeSubscriptionScopeKey, normalizeAssetId } from '../utils'; import { ZERO_ADDRESS } from '../utils/constants'; import { AbstractDataSource } from './AbstractDataSource'; import type { @@ -156,6 +156,10 @@ type SubscriptionData = { chains: ChainId[]; /** Accounts being polled */ accounts: InternalAccount[]; + /** Stable scope key for this subscription (accounts + chains + interval + mode). */ + scopeKey: string; + /** Timestamp (ms) when polling for this scope was (re)started. */ + subscribedAt: number; /** Callback to report asset updates to the controller */ onAssetsUpdate: ( response: DataResponse, @@ -1379,17 +1383,45 @@ export class RpcDataSource extends AbstractDataSource< return; } + const pollInterval = + request.updateInterval ?? + this.#balanceFetcher.getIntervalLength() ?? + DEFAULT_BALANCE_INTERVAL; + const scopeKey = computeSubscriptionScopeKey( + request, + chainsToSubscribe, + pollInterval, + ); + + const existing = this.#activeSubscriptions.get(subscriptionId); + + // Skip a redundant restart when an existing subscription already covers the + // exact same scope and its polling (re)started within one interval. + // Restarting would stop/start polling and trigger an immediate duplicate + // RPC fetch for data we just requested (e.g. a chain flapping between the + // WebSocket source and RPC). `forceUpdate` always restarts (it is excluded + // from the scope key). The existing polling keeps refreshing on its own + // interval, so no data is lost. + if ( + !request.forceUpdate && + existing?.scopeKey === scopeKey && + Date.now() - (existing?.subscribedAt ?? 0) < pollInterval + ) { + log('Skipping redundant re-subscribe (scope unchanged)', { + subscriptionId, + chains: chainsToSubscribe, + }); + return; + } + // Handle subscription update - restart polling for new chains - if (isUpdate) { - const existing = this.#activeSubscriptions.get(subscriptionId); - if (existing) { - log('Updating existing subscription - restarting polling', { - subscriptionId, - existingChains: existing.chains, - newChains: chainsToSubscribe, - }); - // Don't return early - continue to unsubscribe and restart polling - } + if (isUpdate && existing) { + log('Updating existing subscription - restarting polling', { + subscriptionId, + existingChains: existing.chains, + newChains: chainsToSubscribe, + }); + // Don't return early - continue to unsubscribe and restart polling } // Clean up existing subscription (stops old polling) @@ -1454,6 +1486,8 @@ export class RpcDataSource extends AbstractDataSource< detectionPollingTokens, chains: chainsToSubscribe, accounts, + scopeKey, + subscribedAt: Date.now(), onAssetsUpdate: subscriptionRequest.onAssetsUpdate, }); diff --git a/packages/assets-controller/src/data-sources/index.ts b/packages/assets-controller/src/data-sources/index.ts index 47cb7c5734c..1d8321fcc82 100644 --- a/packages/assets-controller/src/data-sources/index.ts +++ b/packages/assets-controller/src/data-sources/index.ts @@ -21,6 +21,14 @@ export { type BackendWebsocketDataSourceAllowedEvents, } from './BackendWebsocketDataSource'; +export { + AccountActivityDataSource, + createAccountActivityDataSource, + type AccountActivityDataSourceOptions, + type AccountActivityDataSourceState, + type AccountActivityDataSourceAllowedEvents, +} from './AccountActivityDataSource'; + export { RpcDataSource, createRpcDataSource, diff --git a/packages/assets-controller/src/utils/index.ts b/packages/assets-controller/src/utils/index.ts index 408f745261d..8d35f306040 100644 --- a/packages/assets-controller/src/utils/index.ts +++ b/packages/assets-controller/src/utils/index.ts @@ -13,3 +13,4 @@ export { buildNativeAssetsFromConstant, buildNativeAssetsFromApi, } from './native-assets'; +export { computeSubscriptionScopeKey } from './subscriptionScope'; diff --git a/packages/assets-controller/src/utils/subscriptionScope.test.ts b/packages/assets-controller/src/utils/subscriptionScope.test.ts new file mode 100644 index 00000000000..476ca8dc0dc --- /dev/null +++ b/packages/assets-controller/src/utils/subscriptionScope.test.ts @@ -0,0 +1,124 @@ +import type { InternalAccount } from '@metamask/keyring-internal-api'; + +import type { ChainId, DataRequest } from '../types'; +import { computeSubscriptionScopeKey } from './subscriptionScope'; + +const CHAIN_MAINNET = 'eip155:1' as ChainId; +const CHAIN_POLYGON = 'eip155:137' as ChainId; + +function createAccount(overrides?: Partial): InternalAccount { + return { + id: 'account-1', + address: '0xAbC0000000000000000000000000000000000001', + options: {}, + methods: [], + type: 'eip155:eoa', + scopes: ['eip155:0'], + metadata: { + name: 'Account 1', + keyring: { type: 'HD Key Tree' }, + importTime: 0, + lastSelected: 0, + }, + ...overrides, + } as InternalAccount; +} + +function createRequest(overrides?: Partial): DataRequest { + return { + accountsWithSupportedChains: [ + { account: createAccount(), supportedChains: [CHAIN_MAINNET] }, + ], + chainIds: [CHAIN_MAINNET], + dataTypes: ['balance'], + ...overrides, + }; +} + +describe('computeSubscriptionScopeKey', () => { + it('produces identical keys for equivalent scopes regardless of ordering', () => { + const requestA = createRequest({ + accountsWithSupportedChains: [ + { + account: createAccount({ id: 'a' }), + supportedChains: [CHAIN_MAINNET, CHAIN_POLYGON], + }, + { account: createAccount({ id: 'b' }), supportedChains: [CHAIN_MAINNET] }, + ], + chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], + }); + const requestB = createRequest({ + accountsWithSupportedChains: [ + { account: createAccount({ id: 'b' }), supportedChains: [CHAIN_MAINNET] }, + { + account: createAccount({ id: 'a' }), + supportedChains: [CHAIN_POLYGON, CHAIN_MAINNET], + }, + ], + chainIds: [CHAIN_POLYGON, CHAIN_MAINNET], + }); + + expect( + computeSubscriptionScopeKey(requestA, [CHAIN_MAINNET, CHAIN_POLYGON], 30000), + ).toBe( + computeSubscriptionScopeKey(requestB, [CHAIN_POLYGON, CHAIN_MAINNET], 30000), + ); + }); + + it('is case-insensitive for account addresses', () => { + const lower = createRequest({ + accountsWithSupportedChains: [ + { + account: createAccount({ + address: '0xabc0000000000000000000000000000000000001', + }), + supportedChains: [CHAIN_MAINNET], + }, + ], + }); + const upper = createRequest({ + accountsWithSupportedChains: [ + { + account: createAccount({ + address: '0xABC0000000000000000000000000000000000001', + }), + supportedChains: [CHAIN_MAINNET], + }, + ], + }); + + expect(computeSubscriptionScopeKey(lower, [CHAIN_MAINNET], 30000)).toBe( + computeSubscriptionScopeKey(upper, [CHAIN_MAINNET], 30000), + ); + }); + + it('differs when the chains differ', () => { + const request = createRequest(); + expect(computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 30000)).not.toBe( + computeSubscriptionScopeKey(request, [CHAIN_POLYGON], 30000), + ); + }); + + it('differs when the poll interval differs', () => { + const request = createRequest(); + expect(computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 30000)).not.toBe( + computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 60000), + ); + }); + + it('differs when customAssetsOnly differs', () => { + const regular = createRequest(); + const customOnly = createRequest({ customAssetsOnly: true }); + expect( + computeSubscriptionScopeKey(regular, [CHAIN_MAINNET], 30000), + ).not.toBe(computeSubscriptionScopeKey(customOnly, [CHAIN_MAINNET], 30000)); + }); + + it('ignores forceUpdate so a forced refresh is not treated as a new scope', () => { + const regular = createRequest(); + const forced = createRequest({ forceUpdate: true }); + expect(computeSubscriptionScopeKey(regular, [CHAIN_MAINNET], 30000)).toBe( + computeSubscriptionScopeKey(forced, [CHAIN_MAINNET], 30000), + ); + }); +}); diff --git a/packages/assets-controller/src/utils/subscriptionScope.ts b/packages/assets-controller/src/utils/subscriptionScope.ts new file mode 100644 index 00000000000..5b1843196b8 --- /dev/null +++ b/packages/assets-controller/src/utils/subscriptionScope.ts @@ -0,0 +1,45 @@ +import type { ChainId, DataRequest } from '../types'; + +/** + * Build a stable, order-independent key describing the effective scope of a + * polling subscription (which accounts, on which chains, at which interval, + * and in which mode). Two subscribe calls that produce the same key would issue + * an identical fetch, so a data source can use this to detect and skip a + * redundant immediate fetch when it is torn down and re-created for the same + * scope (e.g. a chain flapping between the WebSocket source and a polling + * source). + * + * The key intentionally excludes `forceUpdate`: callers that force a refresh + * always want a fresh fetch and should never be treated as redundant. + * + * @param request - The data request for the subscription. + * @param chains - The chains actually being subscribed (may be a subset of + * `request.chainIds`). + * @param pollInterval - The resolved polling interval (ms). + * @returns A deterministic scope key string. + */ +export function computeSubscriptionScopeKey( + request: DataRequest, + chains: ChainId[], + pollInterval: number, +): string { + const accounts = request.accountsWithSupportedChains + .map(({ account, supportedChains }) => { + const scopedChains = [...supportedChains].sort().join(','); + return `${account.id}:${account.address.toLowerCase()}:${scopedChains}`; + }) + .sort() + .join('|'); + + const sortedChains = [...chains].sort().join(','); + const customAssetsOnly = request.customAssetsOnly === true ? '1' : '0'; + const customAssets = [...(request.customAssets ?? [])].sort().join(','); + + return [ + accounts, + sortedChains, + String(pollInterval), + customAssetsOnly, + customAssets, + ].join('#'); +} From 3ddda907175b8cae678b95f10dfd8e5b4986a0eb Mon Sep 17 00:00:00 2001 From: Kriys94 Date: Fri, 17 Jul 2026 13:45:36 +0200 Subject: [PATCH 2/2] fix(assets-controller): remove BackendWebsocket --- packages/assets-controller/CHANGELOG.md | 6 + .../src/AssetsController.test.ts | 4 +- .../assets-controller/src/AssetsController.ts | 69 +- packages/assets-controller/src/README.md | 18 +- .../MockAssetControllerMessenger.ts | 19 +- .../data-sources/AccountActivityDataSource.ts | 14 +- .../AccountsApiDataSource.test.ts | 129 -- .../src/data-sources/AccountsApiDataSource.ts | 42 +- .../BackendWebsocketDataSource.test.ts | 1417 ----------------- .../BackendWebsocketDataSource.ts | 843 ---------- .../src/data-sources/RpcDataSource.test.ts | 75 +- .../src/data-sources/RpcDataSource.ts | 56 +- .../src/data-sources/index.ts | 9 - packages/assets-controller/src/index.ts | 11 - .../CustomAssetGraduationMiddleware.test.ts | 4 +- .../CustomAssetGraduationMiddleware.ts | 2 +- packages/assets-controller/src/utils/index.ts | 1 - .../processAccountActivityBalanceUpdates.ts | 2 +- .../src/utils/subscriptionScope.test.ts | 124 -- .../src/utils/subscriptionScope.ts | 45 - 20 files changed, 77 insertions(+), 2813 deletions(-) delete mode 100644 packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts delete mode 100644 packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts delete mode 100644 packages/assets-controller/src/utils/subscriptionScope.test.ts delete mode 100644 packages/assets-controller/src/utils/subscriptionScope.ts diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index d52bcc6eae6..dc97dc83ade 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -18,6 +18,12 @@ 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)) +- `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 + - `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 ## [11.0.0] diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index f50b1e98f27..10f029f5200 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -744,7 +744,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); @@ -756,7 +756,7 @@ describe('AssetsController', () => { }, }, }, - 'BackendWebsocketDataSource', + 'AccountActivityDataSource', ); expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined(); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 82689ee7a80..70072293038 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -17,8 +17,6 @@ import type { ApiPlatformClient, AccountActivityServiceBalanceUpdatedEvent, AccountActivityServiceStatusChangedEvent, - BackendWebSocketServiceActions, - BackendWebSocketServiceEvents, SupportedCurrency, } from '@metamask/core-backend'; import type { @@ -79,7 +77,6 @@ import type { import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource'; import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource'; import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource'; -import { BackendWebsocketDataSource } from './data-sources/BackendWebsocketDataSource'; import { shouldSkipNativeForCaipChainId } from './data-sources/evm-rpc-services/utils/assets'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource'; import { PriceDataSource } from './data-sources/PriceDataSource'; @@ -339,8 +336,6 @@ type AllowedActions = | SnapControllerGetRunnableSnapsAction | SnapControllerHandleRequestAction | GetPermissions - // BackendWebsocketDataSource - | BackendWebSocketServiceActions // PhishingController | PhishingControllerBulkScanTokensAction // AccountsApiDataSource (Accounts API v6 balances feature flag) @@ -370,8 +365,6 @@ type AllowedEvents = | AccountsControllerAccountBalancesUpdatedEvent | PermissionControllerStateChange | SnapControllerSnapInstalledEvent - // BackendWebsocketDataSource - | BackendWebSocketServiceEvents // AccountActivityService (real-time balance updates + chain status for unified assets) | AccountActivityServiceBalanceUpdatedEvent | AccountActivityServiceStatusChangedEvent; @@ -813,8 +806,6 @@ export class AssetsController extends BaseController< return []; } - readonly #backendWebsocketDataSource: BackendWebsocketDataSource; - readonly #accountActivityDataSource: AccountActivityDataSource; readonly #accountsApiDataSource: AccountsApiDataSource; @@ -827,19 +818,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, @@ -930,13 +927,6 @@ export class AssetsController extends BaseController< } }; - this.#backendWebsocketDataSource = new BackendWebsocketDataSource({ - messenger: this.messenger, - queryApiClient, - onActiveChainsUpdated: this.#onActiveChainsUpdated, - getAssetType: (assetId: Caip19AssetId): 'native' | 'erc20' | 'spl' => - this.#getAssetType(assetId), - }); this.#accountActivityDataSource = new AccountActivityDataSource({ messenger: this.messenger, onActiveChainsUpdated: this.#onActiveChainsUpdated, @@ -2963,7 +2953,6 @@ export class AssetsController extends BaseController< // every source that may have been subscribed. const allSources = [ ...this.#allBalanceDataSources, - this.#accountActivityDataSource, this.#stakedBalanceDataSource, ]; const subscriptionKeys = [...this.#activeSubscriptions.keys()]; @@ -3108,6 +3097,16 @@ export class AssetsController extends BaseController< } } + // AccountActivityDataSource only *reserves* its "up" chains here: they are + // removed from `remainingChains` above so the polling sources (AccountsApi + // /RPC) skip them, but the actual subscription is created separately below + // with all selected accounts and enabled chains so `balanceUpdated` events + // can always be routed to the matching account (chain-agnostic), and it is + // never torn down when it has no claimed chains. + if (source === this.#accountActivityDataSource) { + continue; + } + if (assignedChains.length === 0) { this.#unsubscribeDataSource(source); continue; @@ -3145,10 +3144,12 @@ export class AssetsController extends BaseController< ); // AccountActivityDataSource consumes real-time AccountActivityService events. - // It does not participate in chain-claiming; it is subscribed with all - // selected accounts and enabled chains so it can route `balanceUpdated` - // events to the matching account. Chain up/down is managed internally from - // `statusChanged` (see AccountActivityDataSource). + // It is subscribed with all selected accounts and enabled chains so it can + // route `balanceUpdated` events to the matching account regardless of which + // chain is currently claimed. Chain reservation (so polling sources skip the + // chains it reports "up") is handled in the claiming loop above; chain + // up/down itself is managed internally from `statusChanged` (see + // AccountActivityDataSource). this.#subscribeDataSource( this.#accountActivityDataSource, accounts, @@ -3158,7 +3159,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 @@ -3581,14 +3582,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(); } @@ -3746,12 +3744,12 @@ 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 === 'AccountActivityDataSource'; const enrichmentSources: AssetsDataSource[] = [ @@ -3803,7 +3801,6 @@ export class AssetsController extends BaseController< }); // Destroy instantiated data sources - this.#backendWebsocketDataSource?.destroy?.(); this.#accountActivityDataSource?.destroy?.(); this.#accountsApiDataSource?.destroy?.(); this.#snapDataSource?.destroy?.(); diff --git a/packages/assets-controller/src/README.md b/packages/assets-controller/src/README.md index baebf8606fa..07a24742544 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', ); ``` @@ -797,7 +797,7 @@ const multiChainAssets = await messenger.call( | Order | Data Source | Update Mechanism | Chains | | ----- | -------------------------- | ------------------------ | --------------------- | -| 1 | BackendWebsocketDataSource | Real-time WebSocket push | API-supported EVM | +| 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 | @@ -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 2652534783a..dd2e04f6c9a 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.ts b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts index 6bb227b48f7..4c4fd6db746 100644 --- a/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts @@ -68,9 +68,9 @@ export type AccountActivityDataSourceOptions = { /** * Data source that consumes real-time updates from `AccountActivityService`. * - * Unlike {@link BackendWebsocketDataSource}, which owns its own WebSocket - * channel subscriptions, this data source is a thin consumer of the two - * high-level events that `AccountActivityService` publishes: + * `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 accounts in the @@ -86,10 +86,10 @@ export type AccountActivityDataSourceOptions = { * 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` also - * are never seeded from the accounts API the way `BackendWebsocketDataSource` - * does; they only ever reflect live `statusChanged` events (the service flushes - * all tracked chains as "down" when the WebSocket disconnects, releasing them). + * 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) and is wired into the controller's balance diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index 2ac0cf1f3ce..f01829bf159 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -774,135 +774,6 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('skips the immediate fetch when re-subscribing with an unchanged scope within one poll interval', async () => { - const { controller, apiClient, assetsUpdateHandler } = - await setupController(); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); - - // Re-subscribe with the identical scope (e.g. a chain flapped back from the - // WebSocket source). The immediate fetch is skipped as redundant. - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - expect(assetsUpdateHandler).toHaveBeenCalledTimes(1); - - controller.destroy(); - }); - - it('performs the immediate fetch when re-subscribing with forceUpdate even if the scope is unchanged', async () => { - const { controller, apiClient, assetsUpdateHandler } = - await setupController(); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ forceUpdate: true }), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(2); - - controller.destroy(); - }); - - it('performs the immediate fetch when re-subscribing with a different scope', async () => { - const { controller, apiClient, assetsUpdateHandler } = - await setupController(); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_MAINNET] }), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest({ chainIds: [CHAIN_MAINNET, CHAIN_POLYGON] }), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(2); - - controller.destroy(); - }); - - it('performs the immediate fetch when re-subscribing after the previous fetch has gone stale', async () => { - const { controller, apiClient, assetsUpdateHandler } = - await setupController(); - - const nowSpy = jest.spyOn(Date, 'now'); - const startTime = 1_000_000; - nowSpy.mockReturnValue(startTime); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(1); - - // Advance the clock past the default poll interval (30s) so the previous - // fetch is no longer fresh. - nowSpy.mockReturnValue(startTime + 30_001); - - await controller.subscribe({ - subscriptionId: 'sub-1', - request: createDataRequest(), - isUpdate: false, - onAssetsUpdate: assetsUpdateHandler, - }); - - expect( - apiClient.accounts.fetchV5MultiAccountBalances, - ).toHaveBeenCalledTimes(2); - - nowSpy.mockRestore(); - controller.destroy(); - }); - it('subscribe does nothing when no chains', async () => { const { controller, assetsUpdateHandler } = await setupController(); diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index 73f4cdb3e59..6ac6167b0e4 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -21,11 +21,7 @@ import type { Middleware, AssetsControllerStateInternal, } from '../types'; -import { - computeSubscriptionScopeKey, - fetchWithTimeout, - normalizeAssetId, -} from '../utils'; +import { fetchWithTimeout, normalizeAssetId } from '../utils'; import type { DataSourceState, SubscriptionRequest, @@ -220,15 +216,6 @@ export class AccountsApiDataSource extends AbstractDataSource< /** State accessor from subscriptions (for filtering when tokenDetectionEnabled is false) */ #getAssetsState?: () => AssetsControllerStateInternal; - /** - * Scope key + timestamp of the last fetch, preserved across teardown. Used to - * skip the redundant immediate fetch when the subscription is re-created for - * the same scope while its last fetch is still fresh (e.g. a chain flapping - * between the WebSocket and this source). This source only ever has a single - * subscription, so a single record suffices. - */ - #lastFetch: { scopeKey: string; at: number } | null = null; - constructor(options: AccountsApiDataSourceOptions) { super(CONTROLLER_NAME, { ...defaultState, @@ -765,11 +752,6 @@ export class AccountsApiDataSource extends AbstractDataSource< await this.unsubscribe(subscriptionId); const pollInterval = request.updateInterval ?? this.#pollInterval; - const scopeKey = computeSubscriptionScopeKey( - request, - chainsToSubscribe, - pollInterval, - ); // Create poll function for this subscription const pollFn = async (): Promise => { @@ -779,8 +761,6 @@ export class AccountsApiDataSource extends AbstractDataSource< return; } - this.#lastFetch = { scopeKey, at: Date.now() }; - // Use stored request (which gets updated on account changes) const fetchResponse = await this.fetch({ ...subscription.request, @@ -809,24 +789,6 @@ export class AccountsApiDataSource extends AbstractDataSource< onAssetsUpdate: subscriptionRequest.onAssetsUpdate, }); - // Skip the immediate fetch when this subscription was just re-created for - // the same scope and its last fetch is still fresh (within one poll - // interval). This avoids redundant identical requests when a chain flaps - // between the WebSocket source and this polling source. `forceUpdate` - // always fetches (it is excluded from the scope key). The scheduled poll - // above still refreshes on the next interval tick. - const isRedundantImmediateFetch = - !request.forceUpdate && - this.#lastFetch?.scopeKey === scopeKey && - Date.now() - this.#lastFetch.at < pollInterval; - - if (isRedundantImmediateFetch) { - log('Skipping redundant immediate fetch on re-subscribe', { - subscriptionId, - }); - return; - } - // Initial fetch await pollFn(); } @@ -841,8 +803,6 @@ export class AccountsApiDataSource extends AbstractDataSource< clearInterval(this.#chainsRefreshTimer); } - this.#lastFetch = null; - // Clean up subscriptions super.destroy(); } 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 b2d5af59586..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'; -import type { Caip19AssetId, ChainId, DataRequest } from '../types'; -import { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, -} from './BackendWebsocketDataSource'; -import type { - BackendWebsocketDataSourceAllowedActions, - BackendWebsocketDataSourceAllowedEvents, -} from './BackendWebsocketDataSource'; - -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 b0349b297be..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'; -import { projectLogger, createModuleLogger } from '../logger'; -import type { ChainId, Caip19AssetId, DataResponse } from '../types'; -import { processAccountActivityBalanceUpdates } from '../utils/processAccountActivityBalanceUpdates'; -import { AbstractDataSource } from './AbstractDataSource'; -import type { - DataSourceState, - SubscriptionRequest, -} from './AbstractDataSource'; - -// ============================================================================ -// 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/RpcDataSource.test.ts b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts index 2ae3e8f6b56..c2ff9044d13 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts @@ -1276,49 +1276,11 @@ describe('RpcDataSource', () => { }); }); - it('restarts polling on update when the scope changes', async () => { + it('updates existing subscription when isUpdate true', async () => { const balanceStartSpy = jest.spyOn( BalanceFetcher.prototype, 'startPolling', ); - const secondAccount = createMockInternalAccount({ - id: 'account-2', - address: '0x9999999999999999999999999999999999999999', - }); - await withController(async ({ controller }) => { - await controller.subscribe({ - request: createDataRequest(), - subscriptionId: 'test-sub', - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - balanceStartSpy.mockClear(); - - // Adding a second account changes the subscription scope, so polling is - // restarted for the new set of accounts. - await controller.subscribe({ - request: createDataRequest({ - accounts: [createMockInternalAccount(), secondAccount], - }), - subscriptionId: 'test-sub', - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - expect(balanceStartSpy).toHaveBeenCalledWith( - expect.objectContaining({ accountId: 'account-2' }), - ); - }); - }); - - it('skips restarting polling when re-subscribing with an unchanged scope within one interval', async () => { - const balanceStartSpy = jest.spyOn( - BalanceFetcher.prototype, - 'startPolling', - ); - const balanceStopSpy = jest.spyOn( - BalanceFetcher.prototype, - 'stopPollingByPollingToken', - ); await withController(async ({ controller }) => { await controller.subscribe({ request: createDataRequest(), @@ -1326,46 +1288,13 @@ describe('RpcDataSource', () => { isUpdate: false, onAssetsUpdate: jest.fn(), }); - balanceStartSpy.mockClear(); - balanceStopSpy.mockClear(); - - // Re-subscribing with the identical scope (e.g. a chain flapping back - // from the WebSocket source) is a no-op: polling is neither stopped nor - // restarted, so no redundant immediate RPC fetch is triggered. - await controller.subscribe({ - request: createDataRequest(), - subscriptionId: 'test-sub', - isUpdate: true, - onAssetsUpdate: jest.fn(), - }); - expect(balanceStartSpy).not.toHaveBeenCalled(); - expect(balanceStopSpy).not.toHaveBeenCalled(); - }); - }); - - it('restarts polling when re-subscribing with forceUpdate even if the scope is unchanged', async () => { - const balanceStartSpy = jest.spyOn( - BalanceFetcher.prototype, - 'startPolling', - ); - await withController(async ({ controller }) => { await controller.subscribe({ request: createDataRequest(), subscriptionId: 'test-sub', - isUpdate: false, - onAssetsUpdate: jest.fn(), - }); - balanceStartSpy.mockClear(); - - await controller.subscribe({ - request: createDataRequest({ forceUpdate: true }), - subscriptionId: 'test-sub', isUpdate: true, onAssetsUpdate: jest.fn(), }); - expect(balanceStartSpy).toHaveBeenCalledWith( - expect.objectContaining({ accountId: MOCK_ACCOUNT_ID }), - ); + expect(balanceStartSpy).toHaveBeenCalledTimes(2); }); }); diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.ts b/packages/assets-controller/src/data-sources/RpcDataSource.ts index cd63f841c21..d28af6e3fda 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.ts @@ -35,7 +35,7 @@ import type { DataResponse, Middleware, } from '../types'; -import { computeSubscriptionScopeKey, normalizeAssetId } from '../utils'; +import { normalizeAssetId } from '../utils'; import { ZERO_ADDRESS } from '../utils/constants'; import { AbstractDataSource } from './AbstractDataSource'; import type { @@ -156,10 +156,6 @@ type SubscriptionData = { chains: ChainId[]; /** Accounts being polled */ accounts: InternalAccount[]; - /** Stable scope key for this subscription (accounts + chains + interval + mode). */ - scopeKey: string; - /** Timestamp (ms) when polling for this scope was (re)started. */ - subscribedAt: number; /** Callback to report asset updates to the controller */ onAssetsUpdate: ( response: DataResponse, @@ -1383,45 +1379,17 @@ export class RpcDataSource extends AbstractDataSource< return; } - const pollInterval = - request.updateInterval ?? - this.#balanceFetcher.getIntervalLength() ?? - DEFAULT_BALANCE_INTERVAL; - const scopeKey = computeSubscriptionScopeKey( - request, - chainsToSubscribe, - pollInterval, - ); - - const existing = this.#activeSubscriptions.get(subscriptionId); - - // Skip a redundant restart when an existing subscription already covers the - // exact same scope and its polling (re)started within one interval. - // Restarting would stop/start polling and trigger an immediate duplicate - // RPC fetch for data we just requested (e.g. a chain flapping between the - // WebSocket source and RPC). `forceUpdate` always restarts (it is excluded - // from the scope key). The existing polling keeps refreshing on its own - // interval, so no data is lost. - if ( - !request.forceUpdate && - existing?.scopeKey === scopeKey && - Date.now() - (existing?.subscribedAt ?? 0) < pollInterval - ) { - log('Skipping redundant re-subscribe (scope unchanged)', { - subscriptionId, - chains: chainsToSubscribe, - }); - return; - } - // Handle subscription update - restart polling for new chains - if (isUpdate && existing) { - log('Updating existing subscription - restarting polling', { - subscriptionId, - existingChains: existing.chains, - newChains: chainsToSubscribe, - }); - // Don't return early - continue to unsubscribe and restart polling + if (isUpdate) { + const existing = this.#activeSubscriptions.get(subscriptionId); + if (existing) { + log('Updating existing subscription - restarting polling', { + subscriptionId, + existingChains: existing.chains, + newChains: chainsToSubscribe, + }); + // Don't return early - continue to unsubscribe and restart polling + } } // Clean up existing subscription (stops old polling) @@ -1486,8 +1454,6 @@ export class RpcDataSource extends AbstractDataSource< detectionPollingTokens, chains: chainsToSubscribe, accounts, - scopeKey, - subscribedAt: Date.now(), onAssetsUpdate: subscriptionRequest.onAssetsUpdate, }); diff --git a/packages/assets-controller/src/data-sources/index.ts b/packages/assets-controller/src/data-sources/index.ts index 1d8321fcc82..4889f586beb 100644 --- a/packages/assets-controller/src/data-sources/index.ts +++ b/packages/assets-controller/src/data-sources/index.ts @@ -12,15 +12,6 @@ export { type AccountsApiDataSourceAllowedActions, } from './AccountsApiDataSource'; -export { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, - type BackendWebsocketDataSourceOptions, - type BackendWebsocketDataSourceState, - type BackendWebsocketDataSourceAllowedActions, - type BackendWebsocketDataSourceAllowedEvents, -} from './BackendWebsocketDataSource'; - export { AccountActivityDataSource, createAccountActivityDataSource, diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 4f69c989517..c29a2694228 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -111,17 +111,6 @@ export type { AccountsApiDataSourceState, } from './data-sources'; -// Data sources - BackendWebsocket -export { - BackendWebsocketDataSource, - createBackendWebsocketDataSource, -} from './data-sources'; - -export type { - BackendWebsocketDataSourceOptions, - BackendWebsocketDataSourceState, -} from './data-sources'; - // Data sources - RPC export { RpcDataSource, createRpcDataSource } from './data-sources'; diff --git a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts index 02331161e93..0cf31e4bbff 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 3fec59bcce5..7bb678c71d6 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/index.ts b/packages/assets-controller/src/utils/index.ts index 8d35f306040..408f745261d 100644 --- a/packages/assets-controller/src/utils/index.ts +++ b/packages/assets-controller/src/utils/index.ts @@ -13,4 +13,3 @@ export { buildNativeAssetsFromConstant, buildNativeAssetsFromApi, } from './native-assets'; -export { computeSubscriptionScopeKey } from './subscriptionScope'; diff --git a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts b/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts index 16c3844ce3c..b5a7f25fb17 100644 --- a/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts +++ b/packages/assets-controller/src/utils/processAccountActivityBalanceUpdates.ts @@ -10,7 +10,7 @@ import type { /** * Convert AccountActivityMessage balance updates into a {@link DataResponse} - * for AssetsController (same shape as BackendWebsocketDataSource). + * for AssetsController. * * @param updates - Balance updates from account-activity websocket payload. * @param accountId - Internal account UUID. diff --git a/packages/assets-controller/src/utils/subscriptionScope.test.ts b/packages/assets-controller/src/utils/subscriptionScope.test.ts deleted file mode 100644 index 476ca8dc0dc..00000000000 --- a/packages/assets-controller/src/utils/subscriptionScope.test.ts +++ /dev/null @@ -1,124 +0,0 @@ -import type { InternalAccount } from '@metamask/keyring-internal-api'; - -import type { ChainId, DataRequest } from '../types'; -import { computeSubscriptionScopeKey } from './subscriptionScope'; - -const CHAIN_MAINNET = 'eip155:1' as ChainId; -const CHAIN_POLYGON = 'eip155:137' as ChainId; - -function createAccount(overrides?: Partial): InternalAccount { - return { - id: 'account-1', - address: '0xAbC0000000000000000000000000000000000001', - options: {}, - methods: [], - type: 'eip155:eoa', - scopes: ['eip155:0'], - metadata: { - name: 'Account 1', - keyring: { type: 'HD Key Tree' }, - importTime: 0, - lastSelected: 0, - }, - ...overrides, - } as InternalAccount; -} - -function createRequest(overrides?: Partial): DataRequest { - return { - accountsWithSupportedChains: [ - { account: createAccount(), supportedChains: [CHAIN_MAINNET] }, - ], - chainIds: [CHAIN_MAINNET], - dataTypes: ['balance'], - ...overrides, - }; -} - -describe('computeSubscriptionScopeKey', () => { - it('produces identical keys for equivalent scopes regardless of ordering', () => { - const requestA = createRequest({ - accountsWithSupportedChains: [ - { - account: createAccount({ id: 'a' }), - supportedChains: [CHAIN_MAINNET, CHAIN_POLYGON], - }, - { account: createAccount({ id: 'b' }), supportedChains: [CHAIN_MAINNET] }, - ], - chainIds: [CHAIN_MAINNET, CHAIN_POLYGON], - }); - const requestB = createRequest({ - accountsWithSupportedChains: [ - { account: createAccount({ id: 'b' }), supportedChains: [CHAIN_MAINNET] }, - { - account: createAccount({ id: 'a' }), - supportedChains: [CHAIN_POLYGON, CHAIN_MAINNET], - }, - ], - chainIds: [CHAIN_POLYGON, CHAIN_MAINNET], - }); - - expect( - computeSubscriptionScopeKey(requestA, [CHAIN_MAINNET, CHAIN_POLYGON], 30000), - ).toBe( - computeSubscriptionScopeKey(requestB, [CHAIN_POLYGON, CHAIN_MAINNET], 30000), - ); - }); - - it('is case-insensitive for account addresses', () => { - const lower = createRequest({ - accountsWithSupportedChains: [ - { - account: createAccount({ - address: '0xabc0000000000000000000000000000000000001', - }), - supportedChains: [CHAIN_MAINNET], - }, - ], - }); - const upper = createRequest({ - accountsWithSupportedChains: [ - { - account: createAccount({ - address: '0xABC0000000000000000000000000000000000001', - }), - supportedChains: [CHAIN_MAINNET], - }, - ], - }); - - expect(computeSubscriptionScopeKey(lower, [CHAIN_MAINNET], 30000)).toBe( - computeSubscriptionScopeKey(upper, [CHAIN_MAINNET], 30000), - ); - }); - - it('differs when the chains differ', () => { - const request = createRequest(); - expect(computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 30000)).not.toBe( - computeSubscriptionScopeKey(request, [CHAIN_POLYGON], 30000), - ); - }); - - it('differs when the poll interval differs', () => { - const request = createRequest(); - expect(computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 30000)).not.toBe( - computeSubscriptionScopeKey(request, [CHAIN_MAINNET], 60000), - ); - }); - - it('differs when customAssetsOnly differs', () => { - const regular = createRequest(); - const customOnly = createRequest({ customAssetsOnly: true }); - expect( - computeSubscriptionScopeKey(regular, [CHAIN_MAINNET], 30000), - ).not.toBe(computeSubscriptionScopeKey(customOnly, [CHAIN_MAINNET], 30000)); - }); - - it('ignores forceUpdate so a forced refresh is not treated as a new scope', () => { - const regular = createRequest(); - const forced = createRequest({ forceUpdate: true }); - expect(computeSubscriptionScopeKey(regular, [CHAIN_MAINNET], 30000)).toBe( - computeSubscriptionScopeKey(forced, [CHAIN_MAINNET], 30000), - ); - }); -}); diff --git a/packages/assets-controller/src/utils/subscriptionScope.ts b/packages/assets-controller/src/utils/subscriptionScope.ts deleted file mode 100644 index 5b1843196b8..00000000000 --- a/packages/assets-controller/src/utils/subscriptionScope.ts +++ /dev/null @@ -1,45 +0,0 @@ -import type { ChainId, DataRequest } from '../types'; - -/** - * Build a stable, order-independent key describing the effective scope of a - * polling subscription (which accounts, on which chains, at which interval, - * and in which mode). Two subscribe calls that produce the same key would issue - * an identical fetch, so a data source can use this to detect and skip a - * redundant immediate fetch when it is torn down and re-created for the same - * scope (e.g. a chain flapping between the WebSocket source and a polling - * source). - * - * The key intentionally excludes `forceUpdate`: callers that force a refresh - * always want a fresh fetch and should never be treated as redundant. - * - * @param request - The data request for the subscription. - * @param chains - The chains actually being subscribed (may be a subset of - * `request.chainIds`). - * @param pollInterval - The resolved polling interval (ms). - * @returns A deterministic scope key string. - */ -export function computeSubscriptionScopeKey( - request: DataRequest, - chains: ChainId[], - pollInterval: number, -): string { - const accounts = request.accountsWithSupportedChains - .map(({ account, supportedChains }) => { - const scopedChains = [...supportedChains].sort().join(','); - return `${account.id}:${account.address.toLowerCase()}:${scopedChains}`; - }) - .sort() - .join('|'); - - const sortedChains = [...chains].sort().join(','); - const customAssetsOnly = request.customAssetsOnly === true ? '1' : '0'; - const customAssets = [...(request.customAssets ?? [])].sort().join(','); - - return [ - accounts, - sortedChains, - String(pollInterval), - customAssetsOnly, - customAssets, - ].join('#'); -}