From ba6857202f37fa887cdd24b43a2316794506e054 Mon Sep 17 00:00:00 2001 From: salimtb Date: Wed, 15 Apr 2026 18:14:53 +0200 Subject: [PATCH] fix(assets-controller): improve snap accounts subscription and data source resilience - Re-subscribe with full account group when AccountTree initializes late - Exempt custom assets from spam filtering in TokenDataSource - Fallback to AccountsApiDataSource polling when websocket is disconnected or disabled - Add support for unapproved transaction events in AssetsController - Add changelog entries --- README.md | 1 + packages/assets-controller/CHANGELOG.md | 15 ++ packages/assets-controller/package.json | 1 + .../src/AssetsController.test.ts | 9 + .../assets-controller/src/AssetsController.ts | 199 +++++++++++++++--- .../BackendWebsocketDataSource.test.ts | 20 +- .../BackendWebsocketDataSource.ts | 86 +++++--- .../src/data-sources/RpcDataSource.test.ts | 49 +++++ .../src/data-sources/RpcDataSource.ts | 17 ++ .../src/data-sources/TokenDataSource.ts | 18 +- yarn.lock | 1 + 11 files changed, 351 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 7a518aeaaee..212584ea1be 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,7 @@ linkStyle default opacity:0.5 approval_controller --> base_controller; approval_controller --> messenger; assets_controller --> account_tree_controller; + assets_controller --> accounts_controller; assets_controller --> assets_controllers; assets_controller --> base_controller; assets_controller --> client_controller; diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index d37d46656ae..80662716da7 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added `isOnboarded` option to `AssetsControllerOptions` and `RpcDataSourceConfig` ([#8430](https://github.com/MetaMask/core/pull/8430)) + - When `isOnboarded` returns `false`, `RpcDataSource` skips `fetch` and `subscribe` calls, preventing on-chain RPC calls before onboarding is complete. + - Defaults to `() => true` so existing consumers are unaffected. + ### Changed - Bump `@metamask/account-tree-controller` from `^7.0.0` to `^7.1.0` ([#8472](https://github.com/MetaMask/core/pull/8472)) @@ -21,6 +27,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- `AssetsController` now re-subscribes to all data sources when `AccountTreeController` state changes after initial startup, ensuring snap accounts and their chains are included ([#8430](https://github.com/MetaMask/core/pull/8430)) + - Previously, `#start()` would create subscriptions before snap accounts were available, and its idempotency guard prevented re-subscription when the full account list arrived. + - Added `#handleAccountTreeStateChange()` which forces a full re-subscription and re-fetch when the account tree updates, picking up all accounts including snaps. + - Added `AccountsController:getSelectedAccount` as a fallback in `#getSelectedAccounts()` for when the account tree is not yet initialized. +- `TokenDataSource` no longer filters out user-imported custom assets with the `MIN_TOKEN_OCCURRENCES` spam filter or Blockaid bulk scan ([#8430](https://github.com/MetaMask/core/pull/8430)) + - Tokens present in `customAssets` state now bypass the EVM occurrence threshold and non-EVM Blockaid scan, ensuring manually imported assets always appear. +- `BackendWebsocketDataSource` now properly releases chains to `AccountsApiDataSource` when the websocket is disconnected or disabled ([#8430](https://github.com/MetaMask/core/pull/8430)) + - Previously, `BackendWebsocketDataSource` eagerly claimed all supported chains on initialization regardless of connection state, preventing `AccountsApiDataSource` from polling. + - Chains are now only claimed when the websocket is connected. On disconnect, chains are released so the chain-claiming loop assigns them to `AccountsApiDataSource` for polling fallback. On reconnect, chains are reclaimed. - `AssetsController` no longer silently skips asset fetching on startup for returning users ([#8412](https://github.com/MetaMask/core/pull/8412)) - Previously, `#start()` was called at keyring unlock before `AccountTreeController.init()` had built the account tree, causing `#selectedAccounts` to return an empty array and all subscriptions and fetches to be skipped. `selectedAccountGroupChange` does not fire when the persisted selected group is unchanged, leaving the controller idle. - Now subscribes to `AccountTreeController:stateChange` (the base-controller event guaranteed to fire when `init()` calls `this.update()`), so the controller re-evaluates its active state once accounts are available. diff --git a/packages/assets-controller/package.json b/packages/assets-controller/package.json index 5439af5664e..5e2f0b92ca4 100644 --- a/packages/assets-controller/package.json +++ b/packages/assets-controller/package.json @@ -57,6 +57,7 @@ "@ethersproject/abi": "^5.7.0", "@ethersproject/providers": "^5.7.0", "@metamask/account-tree-controller": "^7.1.0", + "@metamask/accounts-controller": "^37.2.0", "@metamask/assets-controllers": "^104.0.0", "@metamask/base-controller": "^9.1.0", "@metamask/client-controller": "^1.0.1", diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index b7e1e3fed49..b708c4696a7 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -142,6 +142,15 @@ async function withController( namespace: MOCK_ANY_NAMESPACE, }); + // Mock AccountsController + ( + messenger as { + registerActionHandler: (a: string, h: () => unknown) => void; + } + ).registerActionHandler('AccountsController:getSelectedAccount', () => + createMockInternalAccount(), + ); + // Mock AccountTreeController messenger.registerActionHandler( 'AccountTreeController:getAccountsFromSelectedAccountGroup', diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 7214b672441..f7b4062960d 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -3,6 +3,7 @@ import type { AccountTreeControllerSelectedAccountGroupChangeEvent, AccountTreeControllerStateChangeEvent, } from '@metamask/account-tree-controller'; +import type { AccountsControllerGetSelectedAccountAction } from '@metamask/accounts-controller'; import { BaseController } from '@metamask/base-controller'; import type { ControllerGetStateAction, @@ -48,6 +49,8 @@ import type { import type { TransactionControllerIncomingTransactionsReceivedEvent, TransactionControllerTransactionConfirmedEvent, + TransactionControllerUnapprovedTransactionAddedEvent, + TransactionMeta, } from '@metamask/transaction-controller'; import { isCaipChainId, @@ -265,6 +268,7 @@ export type AssetsControllerEvents = type AllowedActions = // AssetsController + | AccountsControllerGetSelectedAccountAction | AccountTreeControllerGetAccountsFromSelectedAccountGroupAction // RpcDataSource | NetworkControllerGetStateAction @@ -288,6 +292,7 @@ type AllowedEvents = | KeyringControllerLockEvent | KeyringControllerUnlockEvent | PreferencesControllerStateChangeEvent + | TransactionControllerUnapprovedTransactionAddedEvent // RpcDataSource, StakedBalanceDataSource | NetworkControllerStateChangeEvent | TransactionControllerTransactionConfirmedEvent @@ -355,6 +360,13 @@ export type AssetsControllerOptions = { priceDataSourceConfig?: PriceDataSourceConfig; /** Optional configuration for StakedBalanceDataSource. */ stakedBalanceDataSourceConfig?: StakedBalanceDataSourceConfig; + /** + * Function returning whether onboarding is complete. When false, + * RPC and staked balance data sources skip fetch and subscribe + * (no on-chain calls until the user has finished onboarding). + * Defaults to () => true. + */ + isOnboarded?: () => boolean; }; // ============================================================================ @@ -620,6 +632,13 @@ export class AssetsController extends BaseController< /** Currently enabled chains from NetworkEnablementController */ #enabledChains: Set = new Set(); + /** + * Snapshot of account IDs that were active when the last subscription/fetch + * cycle ran. Used by #handleAccountTreeStateChange to skip redundant + * re-subscriptions when the account set hasn't actually changed. + */ + #lastKnownAccountIds: ReadonlySet = new Set(); + /** * Get the currently selected accounts from AccountTreeController. * This includes all accounts in the same group as the selected account @@ -627,10 +646,20 @@ export class AssetsController extends BaseController< * * @returns Array of InternalAccount objects from the selected account group. */ - get #selectedAccounts(): InternalAccount[] { - return this.messenger.call( + #getSelectedAccounts(): InternalAccount[] { + const accounts = this.messenger.call( 'AccountTreeController:getAccountsFromSelectedAccountGroup', ); + if (accounts.length > 0) { + return accounts; + } + const selectedAccount = this.messenger.call( + 'AccountsController:getSelectedAccount', + ); + if (selectedAccount) { + return [selectedAccount]; + } + return []; } readonly #backendWebsocketDataSource: BackendWebsocketDataSource; @@ -644,8 +673,7 @@ export class AssetsController extends BaseController< readonly #stakedBalanceDataSource: StakedBalanceDataSource; /** - * All balance data sources (used for unsubscription in #stop so we can clean up - * regardless of current isBasicFunctionality mode). + * 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. * @@ -692,6 +720,7 @@ export class AssetsController extends BaseController< accountsApiDataSourceConfig, priceDataSourceConfig, stakedBalanceDataSourceConfig, + isOnboarded, }: AssetsControllerOptions) { super({ name: CONTROLLER_NAME, @@ -742,6 +771,7 @@ export class AssetsController extends BaseController< messenger: this.messenger, onActiveChainsUpdated: this.#onActiveChainsUpdated, ...rpcConfig, + isOnboarded: rpcConfig.isOnboarded ?? isOnboarded, }); this.#stakedBalanceDataSource = new StakedBalanceDataSource({ messenger: this.messenger, @@ -875,7 +905,7 @@ export class AssetsController extends BaseController< // when init() calls this.update(). #start() is idempotent so // repeated fires are safe. this.messenger.subscribe('AccountTreeController:stateChange', () => { - this.#updateActive(); + this.#handleAccountTreeStateChange(); }); // Subscribe to network enablement changes (only enabledNetworkMap) @@ -905,6 +935,44 @@ export class AssetsController extends BaseController< this.#keyringUnlocked = false; this.#updateActive(); }); + + // Subscribe to unapproved transactions - TXs that need confirmation + // Ensures that balances for the account making transaction are updated (e.g. for gas estimations) + this.messenger.subscribe( + 'TransactionController:unapprovedTransactionAdded', + (transactionMeta: TransactionMeta) => { + this.#onUnapprovedTransactionAdded(transactionMeta); + }, + ); + } + + #onUnapprovedTransactionAdded(transactionMeta: TransactionMeta): void { + const hexChainId = transactionMeta.chainId; + if (!hexChainId) { + return; + } + + const caipChainId = `eip155:${parseInt(hexChainId, 16)}` as ChainId; + const fromAddress = transactionMeta.txParams.from?.toLowerCase(); + if (!fromAddress) { + return; + } + + const matchedAccount = this.#getSelectedAccounts().find( + (account) => account.address.toLowerCase() === fromAddress, + ); + if (!matchedAccount) { + return; + } + + this.getAssets([matchedAccount], { + chainIds: [caipChainId], + forceUpdate: true, + }).catch((error) => { + log('Failed to refresh assets after unapproved transaction added', { + error, + }); + }); } /** @@ -920,6 +988,51 @@ export class AssetsController extends BaseController< } } + /** + * Handle AccountTreeController state changes. + * If already running, re-subscribe only when the set of selected accounts + * has actually changed (e.g. a snap account was added after initial startup). + * This guards against the many tree mutations that don't affect which + * accounts are selected — without this check every tree update would + * trigger a redundant full re-subscribe + forceUpdate fetch. + * If not running yet, delegate to #start() for the normal start flow. + */ + #handleAccountTreeStateChange(): void { + const shouldRun = this.#uiOpen && this.#keyringUnlocked; + if (!shouldRun) { + return; + } + if (this.#activeSubscriptions.size > 0) { + const accounts = this.#getSelectedAccounts(); + const currentIds = new Set(accounts.map((a) => a.id)); + + const accountsChanged = + currentIds.size !== this.#lastKnownAccountIds.size || + [...currentIds].some((id) => !this.#lastKnownAccountIds.has(id)); + + if (!accountsChanged) { + return; + } + + log('Account tree changed with new accounts, re-subscribing', { + previousCount: this.#lastKnownAccountIds.size, + currentCount: currentIds.size, + }); + + this.#lastKnownAccountIds = currentIds; + this.#subscribeAssets(); + this.#ensureNativeBalancesDefaultZero(); + this.getAssets(accounts, { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }).catch((error) => { + log('Failed to fetch assets after tree change', error); + }); + } else { + this.#start(); + } + } + #registerActionHandlers(): void { this.messenger.registerMethodActionHandlers( this, @@ -970,13 +1083,13 @@ export class AssetsController extends BaseController< } // If chains were added and we have selected accounts, do one-time fetch - if (addedChains.length > 0 && this.#selectedAccounts.length > 0) { + if (addedChains.length > 0 && this.#getSelectedAccounts().length > 0) { const addedEnabledChains = addedChains.filter((chain) => this.#enabledChains.has(chain), ); if (addedEnabledChains.length > 0) { log('Fetching balances for newly added chains', { addedEnabledChains }); - this.getAssets(this.#selectedAccounts, { + this.getAssets(this.#getSelectedAccounts(), { chainIds: addedEnabledChains, forceUpdate: true, updateMode: 'merge', @@ -1339,7 +1452,7 @@ export class AssetsController extends BaseController< * @returns Legacy-compatible state for transaction-pay-controller. */ getStateForTransactionPay(): TransactionPayLegacyFormat { - const accounts = this.#selectedAccounts; + const accounts = this.#getSelectedAccounts(); const { nativeAssetIdentifiers } = this.messenger.call( 'NetworkEnablementController:getState', ); @@ -1429,7 +1542,7 @@ export class AssetsController extends BaseController< }); // Fetch data for the newly added custom asset (merge to preserve other chains) - const account = this.#selectedAccounts.find((a) => a.id === accountId); + const account = this.#getSelectedAccounts().find((a) => a.id === accountId); if (account) { const chainId = extractChainId(normalizedAssetId); await this.getAssets([account], { @@ -1549,7 +1662,7 @@ export class AssetsController extends BaseController< return; } - this.getAssets(this.#selectedAccounts, { + this.getAssets(this.#getSelectedAccounts(), { forceUpdate: true, dataTypes: ['price'], assetsForPriceUpdate: Object.values(this.state.assetsBalance).flatMap( @@ -1580,10 +1693,11 @@ export class AssetsController extends BaseController< const existingSubscription = this.#activeSubscriptions.get(subscriptionKey); const isUpdate = existingSubscription !== undefined; + const request = this.#buildDataRequest(accounts, chainIds, { + dataTypes: ['price'], + }); const subscribeReq: SubscriptionRequest = { - request: this.#buildDataRequest(accounts, chainIds, { - dataTypes: ['price'], - }), + request, subscriptionId: subscriptionKey, isUpdate, onAssetsUpdate: (response) => @@ -1675,7 +1789,7 @@ export class AssetsController extends BaseController< * Only adds natives for chains that the account supports (correct accountId ↔ chain mapping). */ #ensureNativeBalancesDefaultZero(): void { - const accounts = this.#selectedAccounts; + const accounts = this.#getSelectedAccounts(); if (accounts.length === 0) { return; } @@ -1798,7 +1912,7 @@ export class AssetsController extends BaseController< })(); // Ensure native tokens have an entry (0 if missing) for chains this account supports - const account = this.#selectedAccounts.find( + const account = this.#getSelectedAccounts().find( (a) => a.id === accountId, ); const nativeAssetIdsForAccount = account @@ -2064,7 +2178,7 @@ export class AssetsController extends BaseController< * subscriptions are already active. */ #start(): void { - const accounts = this.#selectedAccounts; + const accounts = this.#getSelectedAccounts(); const chainIds = [...this.#enabledChains]; if (accounts.length === 0 || chainIds.length === 0) { @@ -2080,6 +2194,7 @@ export class AssetsController extends BaseController< enabledChainCount: chainIds.length, }); + this.#lastKnownAccountIds = new Set(accounts.map((a) => a.id)); this.#subscribeAssets(); this.#ensureNativeBalancesDefaultZero(); this.getAssets(accounts, { @@ -2102,6 +2217,7 @@ export class AssetsController extends BaseController< this.#firstInitFetchReported = false; this.#stateSizeReported = false; + this.#lastKnownAccountIds = new Set(); // Stop price subscription first (uses direct messenger call) this.unsubscribeAssetsPrice(); @@ -2144,22 +2260,20 @@ export class AssetsController extends BaseController< * Subscribe to asset updates for all selected accounts. */ #subscribeAssets(): void { - if (this.#selectedAccounts.length === 0 || this.#enabledChains.size === 0) { + const accounts = this.#getSelectedAccounts(); + const enabledChains = [...this.#enabledChains]; + if (accounts.length === 0 || enabledChains.length === 0) { return; } // Subscribe to balance updates (batched by data source) - this.#subscribeAssetsBalance(this.#selectedAccounts, [ - ...this.#enabledChains, - ]); + this.#subscribeAssetsBalance(accounts, enabledChains); // Subscribe to staked balance updates (separate from regular balance chain-claiming) - this.#subscribeStakedBalance(this.#selectedAccounts, [ - ...this.#enabledChains, - ]); + this.#subscribeStakedBalance(accounts, enabledChains); // Subscribe to price updates for all assets held by selected accounts - this.subscribeAssetsPrice(this.#selectedAccounts, [...this.#enabledChains]); + this.subscribeAssetsPrice(accounts, enabledChains); } /** @@ -2185,8 +2299,7 @@ export class AssetsController extends BaseController< new Set(chainIds), ); const remainingChains = new Set(chainToAccounts.keys()); - - // When basic functionality is on (getter true), use all balance data sources; when off (getter false), RPC only. + // When basic functionality is on, use all balance data sources; when off, RPC only. const balanceDataSources = this.#isBasicFunctionality() ? this.#allBalanceDataSources : [this.#rpcDataSource]; @@ -2379,12 +2492,16 @@ export class AssetsController extends BaseController< > = {}, ): DataRequest { const chainIdSet = new Set(chainIds); + // Force update will pass in additional chains for wildcard account scope + // Enables updates for chains that are not selected (e.g. for unapproved transactions) + const forceUpdateChains = partial.forceUpdate ? chainIds : undefined; const accountsWithSupportedChains: AccountWithSupportedChains[] = accounts.map((account) => ({ account, - supportedChains: this.#getEnabledChainsForAccount(account).filter( - (chain) => chainIdSet.has(chain), - ), + supportedChains: this.#getEnabledChainsForAccount( + account, + forceUpdateChains, + ).filter((chain) => chainIdSet.has(chain)), })); return { accountsWithSupportedChains, @@ -2398,9 +2515,13 @@ export class AssetsController extends BaseController< * Returns the intersection of the account's scopes and the enabled chains. * * @param account - The account to get supported chains for. + * @param additionalChains - Optional extra chains to include for wildcard scopes. * @returns Array of ChainIds that the account supports and are enabled. */ - #getEnabledChainsForAccount(account: InternalAccount): ChainId[] { + #getEnabledChainsForAccount( + account: InternalAccount, + additionalChains?: ChainId[], + ): ChainId[] { // Account scopes are CAIP-2 chain IDs like "eip155:1", "solana:mainnet", "bip122:..." const scopes = account.scopes ?? []; const result: ChainId[] = []; @@ -2415,6 +2536,7 @@ export class AssetsController extends BaseController< // Wildcard scope (e.g., "eip155:0" means all enabled chains in that namespace) if (reference === '0') { + // Add enabled chains for (const chain of this.#enabledChains) { if (isCaipChainId(chain)) { const chainParsed = parseCaipChainId(chain); @@ -2423,6 +2545,15 @@ export class AssetsController extends BaseController< } } } + // Add additional chains + for (const chain of additionalChains ?? []) { + if (isCaipChainId(chain)) { + const chainParsed = parseCaipChainId(chain); + if (chainParsed.namespace === namespace) { + result.push(chain); + } + } + } } else if (namespace === 'eip155' && isStrictHexString(reference)) { // Normalize hex to decimal for EIP155 result.push(`eip155:${parseInt(reference, 16)}` as ChainId); @@ -2439,13 +2570,15 @@ export class AssetsController extends BaseController< // ============================================================================ async #handleAccountGroupChanged(): Promise { - const accounts = this.#selectedAccounts; + const accounts = this.#getSelectedAccounts(); log('Account group changed', { accountCount: accounts.length, accountIds: accounts.map((a) => a.id), }); + this.#lastKnownAccountIds = new Set(accounts.map((a) => a.id)); + // Subscribe and fetch for the new account group this.#subscribeAssets(); if (accounts.length > 0) { @@ -2495,8 +2628,8 @@ export class AssetsController extends BaseController< this.#subscribeAssets(); // Do one-time fetch for newly enabled chains; merge so we keep existing chain balances - if (addedChains.length > 0 && this.#selectedAccounts.length > 0) { - await this.getAssets(this.#selectedAccounts, { + if (addedChains.length > 0 && this.#getSelectedAccounts().length > 0) { + await this.getAssets(this.#getSelectedAccounts(), { chainIds: addedChains, forceUpdate: true, updateMode: 'merge', diff --git a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts index d9f272975d8..87debe544f8 100644 --- a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.test.ts @@ -354,7 +354,11 @@ describe('BackendWebsocketDataSource', () => { triggerConnectionStateChange(WebSocketState.CONNECTED); await new Promise(process.nextTick); - expect(wsSubscribeMock).toHaveBeenCalled(); + // 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(); }); @@ -497,11 +501,12 @@ describe('BackendWebsocketDataSource', () => { controller.destroy(); }); - it('handles WebSocket disconnect by moving subscriptions to pending', async () => { + it('handles WebSocket disconnect by releasing chains and reclaiming on reconnect', async () => { const { controller, wsSubscribeMock, getConnectionInfoMock, + activeChainsUpdateHandler, triggerConnectionStateChange, } = setupController({ initialActiveChains: [CHAIN_MAINNET], @@ -539,10 +544,19 @@ describe('BackendWebsocketDataSource', () => { requestTimeout: 30000, }); + activeChainsUpdateHandler.mockClear(); triggerConnectionStateChange(WebSocketState.CONNECTED); await new Promise(process.nextTick); - expect(wsSubscribeMock).toHaveBeenCalledTimes(2); + // 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(); }); diff --git a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts index 210d93f49d0..1edfafe3183 100644 --- a/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts +++ b/packages/assets-controller/src/data-sources/BackendWebsocketDataSource.ts @@ -227,6 +227,12 @@ export class BackendWebsocketDataSource extends AbstractDataSource< /** 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(); @@ -257,10 +263,17 @@ export class BackendWebsocketDataSource extends AbstractDataSource< async #initializeActiveChains(): Promise { try { const chains = await this.#fetchActiveChains(); - const previous = [...this.state.activeChains]; - this.updateActiveChains(chains, (updatedChains) => - this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), - ); + 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); @@ -273,6 +286,13 @@ export class BackendWebsocketDataSource extends AbstractDataSource< 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); @@ -311,10 +331,12 @@ export class BackendWebsocketDataSource extends AbstractDataSource< 'BackendWebSocketService:connectionStateChanged', (connectionInfo: ConnectionStatePayload) => { if (connectionInfo.state === ('connected' as WebSocketState)) { - this.#processPendingSubscriptions().catch(console.error); + this.#isConnected = true; + this.#handleReconnect(); } else if ( connectionInfo.state === ('disconnected' as WebSocketState) ) { + this.#isConnected = false; this.#handleDisconnect(); } }, @@ -337,19 +359,19 @@ export class BackendWebsocketDataSource extends AbstractDataSource< * Moves all active subscriptions to pending for re-subscription on reconnect. */ #handleDisconnect(): void { - log('WebSocket disconnected, preserving subscriptions for reconnect', { + 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) { - // Mark as update since it was previously active this.#pendingSubscriptions.set(subscriptionId, { ...originalRequest, - isUpdate: false, // Treat as new subscription since server cleared it + isUpdate: false, }); } } @@ -359,30 +381,42 @@ export class BackendWebsocketDataSource extends AbstractDataSource< // 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), + ); + } } /** - * Process any pending subscriptions that were queued while WebSocket was disconnected. + * 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. */ - async #processPendingSubscriptions(): Promise { - if (this.#pendingSubscriptions.size === 0) { - return; - } + #handleReconnect(): void { + log('WebSocket reconnected, reclaiming chains', { + supportedChainCount: this.#supportedChains.length, + pendingSubscriptionCount: this.#pendingSubscriptions.size, + }); - // Process all pending subscriptions - const pendingEntries = Array.from(this.#pendingSubscriptions.entries()); + // 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(); - for (const [subscriptionId, request] of pendingEntries) { - try { - // Remove from pending before processing to avoid infinite loop - this.#pendingSubscriptions.delete(subscriptionId); - await this.subscribe(request); - } catch (error) { - log('Failed to process pending subscription', { - subscriptionId, - error, - }); - } + if (this.#supportedChains.length > 0) { + const previous = [...this.state.activeChains]; + this.updateActiveChains(this.#supportedChains, (updatedChains) => + this.#onActiveChainsUpdated(this.getName(), updatedChains, previous), + ); } } diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts index bfd92db17ba..01bb6876aa0 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts @@ -1752,4 +1752,53 @@ describe('RpcDataSource', () => { }); }); }); + + describe('isOnboarded', () => { + it('returns empty response from fetch when isOnboarded returns false', async () => { + await withController( + { options: { isOnboarded: () => false } }, + async ({ controller }) => { + const response = await controller.fetch(createDataRequest()); + expect(response).toStrictEqual({}); + }, + ); + }); + + it('skips subscribe when isOnboarded returns false', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const detectionStartSpy = jest.spyOn( + TokenDetector.prototype, + 'startPolling', + ); + + await withController( + { options: { isOnboarded: () => false } }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest(), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + + expect(balanceStartSpy).not.toHaveBeenCalled(); + expect(detectionStartSpy).not.toHaveBeenCalled(); + }, + ); + }); + + it('fetches normally when isOnboarded returns true', async () => { + await withController( + { options: { isOnboarded: () => true } }, + async ({ controller }) => { + const response = await controller.fetch(createDataRequest()); + expect(response).toBeDefined(); + expect(response.assetsBalance).toBeDefined(); + }, + ); + }); + }); }); diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.ts b/packages/assets-controller/src/data-sources/RpcDataSource.ts index 0a020d61fb5..8094dd569d4 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.ts @@ -102,6 +102,8 @@ export type RpcDataSourceConfig = { tokenDetectionEnabled?: () => boolean; /** Function returning whether external services are allowed (avoids stale value; default: () => true) */ useExternalService?: () => boolean; + /** Function returning whether onboarding is complete. When false, fetch and subscribe are no-ops. Defaults to () => true. */ + isOnboarded?: () => boolean; timeout?: number; }; @@ -124,6 +126,8 @@ export type RpcDataSourceOptions = { tokenDetectionEnabled?: () => boolean; /** Function returning whether external services are allowed (avoids stale value; default: () => true) */ useExternalService?: () => boolean; + /** Function returning whether onboarding is complete. When false, fetch and subscribe are no-ops. Defaults to () => true. */ + isOnboarded?: () => boolean; }; /** @@ -199,6 +203,8 @@ export class RpcDataSource extends AbstractDataSource< readonly #useExternalService: () => boolean; + readonly #isOnboarded: () => boolean; + /** Currently active chains */ #activeChains: ChainId[] = []; @@ -231,6 +237,7 @@ export class RpcDataSource extends AbstractDataSource< options.tokenDetectionEnabled ?? ((): boolean => true); this.#useExternalService = options.useExternalService ?? ((): boolean => true); + this.#isOnboarded = options.isOnboarded ?? ((): boolean => true); const balanceInterval = options.balanceInterval ?? DEFAULT_BALANCE_INTERVAL; const detectionInterval = @@ -852,6 +859,11 @@ export class RpcDataSource extends AbstractDataSource< } async fetch(request: DataRequest): Promise { + if (!this.#isOnboarded()) { + log('Skipping fetch - onboarding not complete'); + return {}; + } + const response: DataResponse = {}; const chainsToFetch = request.chainIds.filter((chainId) => @@ -1214,6 +1226,11 @@ export class RpcDataSource extends AbstractDataSource< * @param subscriptionRequest - The subscription request details. */ async subscribe(subscriptionRequest: SubscriptionRequest): Promise { + if (!this.#isOnboarded()) { + log('Skipping subscribe - onboarding not complete'); + return; + } + const { request, subscriptionId, isUpdate } = subscriptionRequest; // Use request.chainIds when activeChains is not yet populated (e.g. before diff --git a/packages/assets-controller/src/data-sources/TokenDataSource.ts b/packages/assets-controller/src/data-sources/TokenDataSource.ts index d91c45422dc..9ee8041906c 100644 --- a/packages/assets-controller/src/data-sources/TokenDataSource.ts +++ b/packages/assets-controller/src/data-sources/TokenDataSource.ts @@ -305,9 +305,14 @@ export class TokenDataSource { // Extract response from context const { response } = ctx; - const { assetsInfo: stateMetadata } = ctx.getAssetsState(); + const { assetsInfo: stateMetadata, customAssets } = ctx.getAssetsState(); const assetIdsNeedingMetadata = new Set(); + // Custom assets are user-imported — exempt from spam filtering. + const customAssetIds = new Set( + Object.values(customAssets ?? {}).flat(), + ); + // Always include native asset IDs from NetworkEnablementController for (const nativeAssetId of this.#getNativeAssetIds()) { assetIdsNeedingMetadata.add(nativeAssetId); @@ -409,17 +414,24 @@ export class TokenDataSource { // EVM: require minimum occurrence count to suppress low-signal tokens. // Tokens with no occurrence data (undefined) are treated the same as // zero occurrences and filtered out. + // Custom assets (user-imported) bypass the occurrence filter. const allowedEvmIds = new Set( evmErc20Ids.filter( (id) => + customAssetIds.has(id) || (occurrencesByAssetId.get(id) ?? 0) >= MIN_TOKEN_OCCURRENCES, ), ); // Non-EVM: Blockaid bulk scan. - const allowedNonEvmIds = new Set( - await this.#filterBlockaidSpamTokens(nonEvmTokenIds), + // Custom assets (user-imported) bypass Blockaid filtering. + const nonEvmToScan = nonEvmTokenIds.filter( + (id) => !customAssetIds.has(id), ); + const allowedNonEvmIds = new Set([ + ...nonEvmTokenIds.filter((id) => customAssetIds.has(id)), + ...(await this.#filterBlockaidSpamTokens(nonEvmToScan)), + ]); // Start with every asset the API returned; only remove those that // fail their respective filter (EVM occurrences / non-EVM Blockaid). diff --git a/yarn.lock b/yarn.lock index 78e6a123f84..8ffb0b03270 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2771,6 +2771,7 @@ __metadata: "@ethersproject/abi": "npm:^5.7.0" "@ethersproject/providers": "npm:^5.7.0" "@metamask/account-tree-controller": "npm:^7.1.0" + "@metamask/accounts-controller": "npm:^37.2.0" "@metamask/assets-controllers": "npm:^104.0.0" "@metamask/auto-changelog": "npm:^6.0.0" "@metamask/base-controller": "npm:^9.1.0"