diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 9aa1bbe851f..20791a8c40c 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add an optional `customAssets` option to `AssetsController.getAssets`, overriding which pinned assets are attached to the fetch request (sent to the Accounts API as `includeAssetIds` and to the RPC fallback). When omitted, every asset pinned by the requested accounts is attached, as before ([#9600](https://github.com/MetaMask/core/pull/9600)) + ### Changed - Bump `@metamask/phishing-controller` from `^17.3.0` to `^17.3.1` ([#9746](https://github.com/MetaMask/core/pull/9746)) @@ -51,6 +55,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Removed +- **BREAKING:** Remove `CustomAssetGraduationMiddleware` (and `CustomAssetGraduationMiddlewareOptions`) from the public API. User-pinned custom assets are now treated as "display no matter what" and are no longer auto-removed when a balance is detected ([#9600](https://github.com/MetaMask/core/pull/9600)) +- **BREAKING:** Remove the `customAssetsOnly` field from `DataRequest`. Whether an RPC poll is asset-scoped is now decided inside `RpcDataSource` from the claimed `customAssets` on the subscription request: chains outside its regular assignment get a poll scoped to an explicit `assetIds` list instead of a controller-driven request flag ([#9600](https://github.com/MetaMask/core/pull/9600)) - **BREAKING:** Remove `BackendWebsocketDataSource` and its factory/types (`BackendWebsocketDataSource`, `createBackendWebsocketDataSource`, `BackendWebsocketDataSourceOptions`, `BackendWebsocketDataSourceState`). Real-time balance updates and per-chain status are now consumed from `AccountActivityService` via `AccountActivityDataSource`, which manages the WebSocket connection and subscriptions. Consumers no longer need to delegate `BackendWebSocketService` actions/events to the `AssetsController` messenger ([#9517](https://github.com/MetaMask/core/pull/9517)) ## [11.3.1] diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index a0476a5e145..cd352e0771a 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -18,9 +18,12 @@ import type { AssetsControllerMessenger, AssetsControllerState, } from './AssetsController.js'; +import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource.js'; import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource.js'; +import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource.js'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource.js'; import { PriceDataSource } from './data-sources/PriceDataSource.js'; +import { RpcDataSource } from './data-sources/RpcDataSource.js'; import { TokenDataSource } from './data-sources/TokenDataSource.js'; import { buildDefaultAssetsInfo } from './defaults.js'; import type { Assets3346MigrationState } from './migrations/healAssetsInfoMetadata.js'; @@ -671,6 +674,42 @@ describe('AssetsController', () => { }); }); + it('fetches only the newly added asset instead of every pinned asset', async () => { + // Use a valid checksummed address (DAI token address) + const secondAssetId = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware = jest.fn(async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; + + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + + capturedCustomAssets.length = 0; + await controller.addCustomAsset(MOCK_ACCOUNT_ID, secondAssetId); + }); + + // The fetch triggered by adding the second pin must not re-request the + // first pin — the subscription refresh covers it on the next poll. + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toStrictEqual([secondAssetId]); + } + + middlewareGetter.mockRestore(); + }); + it('does not overwrite an existing balance when re-adding a custom asset', async () => { await withController( { @@ -731,99 +770,6 @@ describe('AssetsController', () => { }); }); - describe('custom asset graduation', () => { - const SOLANA_ASSET_ID = - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; - - it('graduates an EVM custom asset when AccountsApiDataSource reports a balance for it', async () => { - await withController(async ({ controller }) => { - await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); - expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain( - MOCK_ASSET_ID, - ); - - await controller.handleAssetsUpdate( - { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [MOCK_ASSET_ID]: { amount: '1000000' }, - }, - }, - }, - 'AccountsApiDataSource', - ); - - expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined(); - }); - }); - - 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); - - await controller.handleAssetsUpdate( - { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [MOCK_ASSET_ID]: { amount: '1000000' }, - }, - }, - }, - 'AccountActivityDataSource', - ); - - expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined(); - }); - }); - - it('does not graduate when RpcDataSource reports a balance for a custom asset', async () => { - await withController(async ({ controller }) => { - await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); - - await controller.handleAssetsUpdate( - { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [MOCK_ASSET_ID]: { amount: '1000000' }, - }, - }, - }, - 'RpcDataSource', - ); - - expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain( - MOCK_ASSET_ID, - ); - }); - }); - - it('does not graduate a non-EVM (Solana) custom asset', async () => { - await withController( - { - state: { - customAssets: { [MOCK_ACCOUNT_ID]: [SOLANA_ASSET_ID] }, - }, - }, - async ({ controller }) => { - await controller.handleAssetsUpdate( - { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [SOLANA_ASSET_ID]: { amount: '1000000' }, - }, - }, - }, - 'AccountsApiDataSource', - ); - - expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain( - SOLANA_ASSET_ID, - ); - }, - ); - }); - }); - describe('getCustomAssets', () => { it('returns empty array for account with no custom assets', async () => { await withController(({ controller }) => { @@ -1431,6 +1377,256 @@ describe('AssetsController', () => { ); }); + it('forwards user-pinned custom assets to the Accounts API v6 endpoint as includeAssetIds', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.addCustomAsset(MOCK_ACCOUNT_ID, customToken); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + includeAssetIds: expect.arrayContaining([customToken]), + }), + expect.anything(), + ); + }, + ); + }); + + it('scopes the custom assets on the request to the requested chains', async () => { + const mainnetToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const polygonToken = + 'eip155:137/erc20:0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174' as Caip19AssetId; + + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware = jest.fn(async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; + + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, mainnetToken); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, polygonToken); + + capturedCustomAssets.length = 0; + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + }); + + // Pins on chains outside the request are dropped when the request is + // built — every data source would only ignore them at fetch time. + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toStrictEqual([mainnetToken]); + } + + middlewareGetter.mockRestore(); + }); + + it('uses the customAssets option instead of state-pinned assets when provided', async () => { + const pinnedToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const requestedToken = + 'eip155:1/erc20:0x6B175474E89094C44Da98b954EedeAC495271d0F' as Caip19AssetId; + + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware = jest.fn(async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; + + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, pinnedToken); + await controller.addCustomAsset(MOCK_ACCOUNT_ID, requestedToken); + + capturedCustomAssets.length = 0; + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + customAssets: [requestedToken], + }); + }); + + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toStrictEqual([requestedToken]); + } + + middlewareGetter.mockRestore(); + }); + + it('scopes the customAssets option to the requested chains and drops invalid IDs', async () => { + const polygonToken = + 'eip155:137/erc20:0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174' as Caip19AssetId; + + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware = jest.fn(async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; + + await withController(async ({ controller }) => { + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + customAssets: [polygonToken, 'not-a-caip-id' as Caip19AssetId], + }); + }); + + // The off-chain pin and the unparseable ID are both dropped, leaving no + // custom assets on the request. + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toBeUndefined(); + } + + middlewareGetter.mockRestore(); + }); + + it('normalizes asset IDs passed via the customAssets option', async () => { + const capturedCustomAssets: (Caip19AssetId[] | undefined)[] = []; + const accountsApiMiddleware = jest.fn(async (ctx, next) => { + capturedCustomAssets.push(ctx.request.customAssets); + return next(ctx); + }); + const middlewareGetter = jest + .spyOn( + AccountsApiDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(accountsApiMiddleware) as unknown as jest.SpyInstance; + + await withController(async ({ controller }) => { + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + customAssets: [MOCK_ASSET_ID_LOWERCASE], + }); + }); + + expect(capturedCustomAssets.length).toBeGreaterThan(0); + for (const customAssets of capturedCustomAssets) { + expect(customAssets).toStrictEqual([MOCK_ASSET_ID]); + } + + middlewareGetter.mockRestore(); + }); + + it('forwards user-hidden assets to the Accounts API v6 endpoint as excludeAssetIds', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + const hiddenToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }, + async ({ controller }) => { + await flushPromises(); + + controller.hideAsset(hiddenToken); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalledWith( + expect.any(Array), + expect.objectContaining({ + excludeAssetIds: expect.arrayContaining([hiddenToken]), + }), + expect.anything(), + ); + }, + ); + }); + describe('pipeline splitting', () => { it('returns from getAssets before background pipelines complete', async () => { // Spy on handleAssetsUpdate to count how many times state is written. @@ -1508,6 +1704,76 @@ describe('AssetsController', () => { ); }); + it('routes chains carrying unprocessed pinned assets (unprocessedCustomAssets) to the slow-pipeline RPC fetch', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [customToken], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances: jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }), + }, + } as unknown as ApiPlatformClient; + + const rpcRequestChainIds: ChainId[][] = []; + const rpcMiddleware = jest.fn(async (ctx, next) => { + rpcRequestChainIds.push(ctx.request.chainIds); + return next(ctx); + }); + const rpcMiddlewareGetter = jest + .spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) + .mockReturnValue(rpcMiddleware) as unknown as jest.SpyInstance; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.addCustomAsset(MOCK_ACCOUNT_ID, customToken); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + // Slow pipeline is fire-and-forget; let it run. + await flushPromises(); + }, + ); + + // The chain of the unresolved pin (eip155:1) — a chain AccountsApi + // handled and did NOT flag as errored — is still routed to RPC in the + // slow pipeline so the pin gets fetched. + expect(rpcMiddleware).toHaveBeenCalled(); + expect( + rpcRequestChainIds.some((chains) => chains.includes('eip155:1')), + ).toBe(true); + + rpcMiddlewareGetter.mockRestore(); + }); + it('does not run token or price middleware in getAssets pipelines when isBasicFunctionality is false', async () => { const tokenMiddlewareGetter = jest.spyOn( TokenDataSource.prototype, @@ -1926,6 +2192,113 @@ describe('AssetsController', () => { tokenMiddlewareGetter.mockRestore(); priceMiddlewareGetter.mockRestore(); }); + + it('falls back to RPC for chains a subscription update flagged as errored (e.g. unprocessedNetworks)', async () => { + const rpcMiddlewareGetter = jest.spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: ['eip155:1'], + dataTypes: ['balance'], + }; + + await withController(async ({ controller }) => { + rpcMiddlewareGetter.mockClear(); + + await controller.handleAssetsUpdate( + { + assetsBalance: {}, + errors: { 'eip155:1': 'Unprocessed networks' }, + }, + 'AccountsApiDataSource', + request, + ); + }); + + // The RpcFallbackMiddleware pulls the RPC data source middleware only when + // there are errored chains to recover. + expect(rpcMiddlewareGetter).toHaveBeenCalled(); + + rpcMiddlewareGetter.mockRestore(); + }); + + it('falls back to RPC for pinned assets a subscription update reported as unprocessed (unprocessedCustomAssets)', async () => { + const rpcMiddlewareGetter = jest.spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: ['eip155:1'], + dataTypes: ['balance'], + customAssets: [customToken], + }; + + await withController(async ({ controller }) => { + rpcMiddlewareGetter.mockClear(); + + await controller.handleAssetsUpdate( + { + assetsBalance: {}, + unprocessedCustomAssets: [customToken], + }, + 'AccountsApiDataSource', + request, + ); + }); + + // The asset-axis signal also pulls the RPC data source middleware for an + // asset-scoped recovery. + expect(rpcMiddlewareGetter).toHaveBeenCalled(); + + rpcMiddlewareGetter.mockRestore(); + }); + + it('does not run the RPC fallback when a subscription update has no errored chains', async () => { + const rpcMiddlewareGetter = jest.spyOn( + RpcDataSource.prototype, + 'assetsMiddleware', + // @ts-expect-error -- Jest supports `get` for accessor spies; `Spyable` typings omit prototype getters. + 'get', + ) as unknown as jest.SpyInstance; + + const request: DataRequest = { + accountsWithSupportedChains: [], + chainIds: ['eip155:1'], + dataTypes: ['balance'], + }; + + await withController(async ({ controller }) => { + rpcMiddlewareGetter.mockClear(); + + await controller.handleAssetsUpdate( + { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_NATIVE_ASSET_ID]: { amount: '1' }, + }, + }, + }, + 'AccountsApiDataSource', + request, + ); + }); + + expect(rpcMiddlewareGetter).not.toHaveBeenCalled(); + + rpcMiddlewareGetter.mockRestore(); + }); }); describe('getAssetsBalance', () => { @@ -2086,6 +2459,55 @@ describe('AssetsController', () => { }); }); + describe('two-axis subscription handoff (chains + custom assets)', () => { + it('claims pinned assets on account-activity-claimed chains instead of letting them fall through', async () => { + // Account activity claims eip155:1; its stream covers pins, so they + // stay with its subscription instead of falling through to a poller. + jest + .spyOn(AccountActivityDataSource.prototype, 'getActiveChainsSync') + .mockReturnValue(['eip155:1' as ChainId]); + const wsSubscribeSpy = jest + .spyOn(AccountActivityDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + const rpcSubscribeSpy = jest + .spyOn(RpcDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + + await withController(async ({ controller }) => { + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + + const wsRequest = wsSubscribeSpy.mock.calls.at(-1)?.[0].request; + expect(wsRequest?.chainIds).toStrictEqual(['eip155:1']); + expect(wsRequest?.customAssets).toStrictEqual([MOCK_ASSET_ID]); + + // Nothing was left for lower-priority sources to claim. + expect(rpcSubscribeSpy).not.toHaveBeenCalled(); + }); + }); + + it('does not create an RPC subscription for pinned assets no source can claim', async () => { + // Account activity is not active on the pin's chain, so the pin falls + // through the whole handoff... + jest + .spyOn(AccountActivityDataSource.prototype, 'getActiveChainsSync') + .mockReturnValue([]); + jest + .spyOn(AccountActivityDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + const rpcSubscribeSpy = jest + .spyOn(RpcDataSource.prototype, 'subscribe') + .mockResolvedValue(undefined); + + await withController(async ({ controller }) => { + // ...and RPC has no provider for the chain (no networks configured in + // the mocked NetworkController), so its real claim returns nothing. + await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID); + + expect(rpcSubscribeSpy).not.toHaveBeenCalled(); + }); + }); + }); + describe('handleAssetsUpdate - state updates', () => { it('updates state with balance data', async () => { await withController(async ({ controller }) => { diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 7b4dd199139..46b8f902b99 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -105,7 +105,6 @@ import { } from './defaults.js'; import { AssetsDataSourceError } from './errors.js'; import { projectLogger, createModuleLogger } from './logger.js'; -import { CustomAssetGraduationMiddleware } from './middlewares/CustomAssetGraduationMiddleware.js'; import { DetectionMiddleware } from './middlewares/DetectionMiddleware.js'; import { createParallelBalanceMiddleware, @@ -139,7 +138,6 @@ import type { Asset, } from './types.js'; import { ZERO_ADDRESS } from './utils/constants.js'; -import { pickRpcCustomAssetsSupplement } from './utils/customAssetsRpcSupplement.js'; import { normalizeAmountString, normalizeAssetId, @@ -813,8 +811,6 @@ export class AssetsController extends BaseController< readonly #detectionMiddleware: DetectionMiddleware; - readonly #customAssetGraduationMiddleware: CustomAssetGraduationMiddleware; - readonly #rpcFallbackMiddleware: RpcFallbackMiddleware; readonly #tokenDataSource: TokenDataSource; @@ -952,19 +948,6 @@ export class AssetsController extends BaseController< ...priceDataSourceConfig, }); this.#detectionMiddleware = new DetectionMiddleware(); - this.#customAssetGraduationMiddleware = new CustomAssetGraduationMiddleware( - { - getSelectedAccountId: (): AccountId | undefined => { - try { - return this.#getSelectedAccounts()[0]?.id; - } catch { - return undefined; - } - }, - removeCustomAsset: (accountId, assetId): void => - this.removeCustomAsset(accountId, assetId), - }, - ); this.#rpcFallbackMiddleware = new RpcFallbackMiddleware({ rpcDataSource: this.#rpcDataSource, }); @@ -1596,6 +1579,12 @@ export class AssetsController extends BaseController< assetsForPriceUpdate?: Caip19AssetId[]; /** When set to `'merge'`, fetch result is merged with existing state instead of replacing. Use for partial fetches (e.g. newly added chains). */ updateMode?: AssetsUpdateMode; + /** + * Pinned assets to attach to this fetch instead of every pin in state. + * Use for targeted fetches (e.g. a newly added token). Entries outside + * the requested chains are dropped. + */ + customAssets?: Caip19AssetId[]; }, ): Promise>> { const chainIds = options?.chainIds ?? [...this.#enabledChains]; @@ -1606,12 +1595,29 @@ export class AssetsController extends BaseController< return this.#getAssetsFromState(accounts, chainIds, assetTypes); } - // Collect custom assets for all requested accounts - const customAssets: Caip19AssetId[] = []; - for (const account of accounts) { - const accountCustomAssets = this.getCustomAssets(account.id); - customAssets.push(...accountCustomAssets); + // Pinned assets for this fetch: the caller's override when provided, + // otherwise every pin of the requested accounts — always scoped to the + // requested chains. + const requestedChains = new Set(chainIds); + const candidateCustomAssets = + options?.customAssets ?? + accounts.flatMap((account) => this.getCustomAssets(account.id)); + const customAssetsSet = new Set(); + for (const assetId of candidateCustomAssets) { + try { + const normalizedAssetId = normalizeAssetId(assetId); + if ( + requestedChains.has(parseCaipAssetType(normalizedAssetId).chainId) + ) { + customAssetsSet.add(normalizedAssetId); + } + } catch { + // Skip unparseable asset IDs + } } + const customAssets = [...customAssetsSet]; + + const hiddenAssets = this.#getHiddenAssetIds(); if (options?.forceUpdate) { // Pipeline spans only on unlock/first-init fetch; later forceUpdates pass @@ -1624,6 +1630,7 @@ export class AssetsController extends BaseController< assetTypes, dataTypes, customAssets: customAssets.length > 0 ? customAssets : undefined, + excludeAssetIds: hiddenAssets.length > 0 ? hiddenAssets : undefined, forceUpdate: true, assetsForPriceUpdate: options?.assetsForPriceUpdate, }); @@ -1632,6 +1639,9 @@ export class AssetsController extends BaseController< // Snap and RPC are excluded here due to their latency (snap triggers account // creation, RPC is slow on many chains). Results are committed to state // immediately so the UI can display balances without waiting for them. + // Errored chains and unresolved pins (`unprocessedCustomAssets`) are + // recovered by RPC in the slow pipeline (`#getSlowPipelineChainIds`), + // keeping RPC off the fast path. // // Fast/slow pipelines use merge so partial API snapshots cannot wipe // tokens missing from the response (e.g. USDC when only native balance @@ -1642,11 +1652,6 @@ export class AssetsController extends BaseController< this.#accountsApiDataSource, this.#stakedBalanceDataSource, ]), - // Graduation must run BEFORE the RPC fallback so it only sees - // AccountsApi/Websocket balances. RPC intentionally carries - // custom assets and must never trigger graduation. - this.#customAssetGraduationMiddleware, - this.#rpcFallbackMiddleware, this.#detectionMiddleware, createParallelMiddleware([ this.#tokenDataSource, @@ -2127,7 +2132,8 @@ export class AssetsController extends BaseController< balances[accountId][normalizedAssetId] ??= { amount: '0' }; }); - // Fetch data for the newly added custom asset (merge to preserve other chains) + // Fetch only the new pin (merge preserves other chains); the account's + // other pins are already covered by the subscription. const account = this.#getSelectedAccounts().find((a) => a.id === accountId); if (account) { const chainId = extractChainId(normalizedAssetId); @@ -2137,11 +2143,11 @@ export class AssetsController extends BaseController< assetTypes: ['fungible'], forceUpdate: true, updateMode: 'merge', + customAssets: [normalizedAssetId], }); } - // Re-evaluate subscriptions so the supplemental RPC poll picks up the - // new customAsset on chains another data source already owns. + // Re-evaluate subscriptions so polls pick up the new pin. this.#subscribeAssets(); } @@ -2169,8 +2175,7 @@ export class AssetsController extends BaseController< } }); - // Re-evaluate subscriptions so the supplemental RPC poll for that chain - // is torn down when no more customAssets remain there. + // Re-evaluate subscriptions so polls drop the removed pin. this.#subscribeAssets(); } @@ -2206,6 +2211,9 @@ export class AssetsController extends BaseController< } state.assetPreferences[normalizedAssetId].hidden = true; }); + + // Re-evaluate subscriptions so polls exclude the newly hidden asset. + this.#subscribeAssets(); } /** @@ -2227,6 +2235,27 @@ export class AssetsController extends BaseController< } } }); + + // Re-evaluate subscriptions so polls stop excluding the asset. + this.#subscribeAssets(); + } + + /** + * Collect globally hidden asset IDs (from `assetPreferences`), forwarded on + * data requests as `excludeAssetIds`. + * + * @returns The CAIP-19 asset IDs the user has hidden. + */ + #getHiddenAssetIds(): Caip19AssetId[] { + const hidden: Caip19AssetId[] = []; + for (const [assetId, prefs] of Object.entries( + this.state.assetPreferences, + )) { + if (prefs.hidden) { + hidden.push(assetId as Caip19AssetId); + } + } + return hidden; } // ============================================================================ @@ -2524,10 +2553,21 @@ export class AssetsController extends BaseController< this.#accountsApiDataSource.getActiveChainsSync(), ); + // Chains whose pins went unresolved (`unprocessedCustomAssets`): route + // them to the slow pipeline so RPC fetches the pins. + const unprocessedCustomAssetChains = new Set( + (fastResponse.unprocessedCustomAssets ?? []).map( + (assetId) => assetId.split('/')[0] as ChainId, + ), + ); + return chainIds.filter((chainId) => { if (fastResponse.errors?.[chainId]) { return true; } + if (unprocessedCustomAssetChains.has(chainId)) { + return true; + } if (!accountsApiChains.has(chainId)) { return true; } @@ -3113,16 +3153,10 @@ export class AssetsController extends BaseController< if (!subscriptionKey.startsWith('ds:')) { continue; } - // Subscription keys take the form `ds:` for the regular - // subscription or `ds::` for supplemental - // subscriptions (e.g. `ds:RpcDataSource:custom`). Split on `:` and - // pick the source-name segment so both shapes resolve correctly. + // Subscription keys take the form `ds:`. const [, sourceId] = subscriptionKey.split(':'); const source = allSources.find((ds) => ds.getName() === sourceId); if (source) { - // Unsubscribe by the actual key — `#unsubscribeDataSource` only - // knows the regular `ds:` shape and would miss - // supplemental subscriptions, leaking their polling timers. this.#unsubscribeBySubscriptionKey(source, subscriptionKey); } } @@ -3168,10 +3202,12 @@ export class AssetsController extends BaseController< * Strategy to minimize data source calls: * 1. Collect all chains to subscribe based on enabled networks * 2. Map chains to accounts based on their scopes - * 3. Split by data source (ordered by priority) - each data source gets ONE subscription + * 3. Split by data source (priority order) - each source gets ONE + * subscription, claiming chains AND pinned assets (`claimCustomAssets`); + * unclaimed assets fall through to lower-priority sources. * * This ensures we make minimal subscriptions to each data source while covering - * all accounts and chains. + * all accounts, chains, and pinned assets. * * @param accounts - Accounts to subscribe balance updates for. * @param chainIds - Chain IDs to subscribe for. @@ -3185,6 +3221,20 @@ export class AssetsController extends BaseController< new Set(chainIds), ); const remainingChains = new Set(chainToAccounts.keys()); + // Pins on enabled chains, offered to the sources in priority order; pins + // on disabled chains are not worth polling. + const remainingCustomAssets = new Set(); + for (const account of accounts) { + for (const assetId of this.getCustomAssets(account.id)) { + try { + if (remainingChains.has(parseCaipAssetType(assetId).chainId)) { + remainingCustomAssets.add(assetId); + } + } catch { + // Skip unparseable asset IDs + } + } + } // When basic functionality is on, use all balance data sources; when off, // RPC only. const isBasicFunctionality = this.#isBasicFunctionality(); @@ -3192,8 +3242,6 @@ export class AssetsController extends BaseController< ? this.#allBalanceDataSources : [this.#rpcDataSource]; - let rpcAssignedChains: Set = new Set(); - for (const source of balanceDataSources) { const availableChains = new Set(source.getActiveChainsSync()); const assignedChains: ChainId[] = []; @@ -3205,15 +3253,20 @@ export class AssetsController extends BaseController< } } - if (assignedChains.length === 0) { - this.#unsubscribeDataSource(source); - continue; + const claimedAssets = source.claimCustomAssets( + [...remainingCustomAssets], + assignedChains, + ); + for (const assetId of claimedAssets) { + remainingCustomAssets.delete(assetId); } - if (source === this.#rpcDataSource) { - rpcAssignedChains = new Set(assignedChains); + if (assignedChains.length === 0 && claimedAssets.length === 0) { + this.#unsubscribeDataSource(source); + continue; } + const claimedAssetsSet = new Set(claimedAssets); const seenIds = new Set(); const accountsForSource = assignedChains .flatMap((chainId) => chainToAccounts.get(chainId) ?? []) @@ -3224,70 +3277,37 @@ export class AssetsController extends BaseController< seenIds.add(account.id); return true; }); + // Owners of claimed assets must be on the subscription even when none + // of their chains were assigned to this source (asset-only coverage). + if (claimedAssetsSet.size > 0) { + for (const account of accounts) { + if (seenIds.has(account.id)) { + continue; + } + if ( + this.getCustomAssets(account.id).some((assetId) => + claimedAssetsSet.has(assetId), + ) + ) { + seenIds.add(account.id); + accountsForSource.push(account); + } + } + } if (accountsForSource.length > 0) { - this.#subscribeDataSource(source, accountsForSource, assignedChains); + this.#subscribeDataSource(source, accountsForSource, assignedChains, { + customAssets: claimedAssets, + }); + } else { + this.#unsubscribeDataSource(source); } } - // Supplemental RPC subscription for customAssets on chains another data - // source claimed during regular handoff. RPC is the sole balance fetcher - // for customAssets, so we must always poll them — even when (e.g.) - // AccountsApi is already covering the chain for normal balances. The - // supplemental subscription runs in `customAssetsOnly` mode so it does - // NOT double-poll the regular tracked balances. - this.#subscribeRpcCustomAssetsSupplement( - accounts, - chainToAccounts, - rpcAssignedChains, - ); - } - - /** - * Guarantee that customAssets are **always** polled by RPC, even when - * 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 - * distinct subscription key (`ds:RpcDataSource:custom`) that does not - * interfere with the regular RPC subscription. - * - * @param accounts - Accounts to consider for customAssets. - * @param chainToAccounts - Map of chain → accounts (built by caller). - * @param rpcAssignedChains - Chains RPC was assigned in the regular handoff. - */ - #subscribeRpcCustomAssetsSupplement( - accounts: InternalAccount[], - chainToAccounts: Map, - rpcAssignedChains: Set, - ): void { - const rpc = this.#rpcDataSource; - const supplementalKey = `ds:${rpc.getName()}:custom`; - - const decision = pickRpcCustomAssetsSupplement({ - accountIds: accounts.map((account) => account.id), - customAssetsByAccount: this.state.customAssets, - rpcAssignedChains, - rpcAvailableChains: new Set(rpc.getActiveChainsSync()), - enabledChains: new Set(chainToAccounts.keys()), - }); - - if (decision.chains.length === 0) { - this.#unsubscribeBySubscriptionKey(rpc, supplementalKey); - return; - } - - const supplementalAccounts = accounts.filter((account) => - decision.accountIds.has(account.id), - ); - if (supplementalAccounts.length === 0) { - this.#unsubscribeBySubscriptionKey(rpc, supplementalKey); - return; + if (remainingCustomAssets.size > 0) { + log('Custom assets unclaimed by any data source', { + assetIds: [...remainingCustomAssets], + }); } - - this.#subscribeDataSource(rpc, supplementalAccounts, decision.chains, { - subscriptionKey: supplementalKey, - customAssetsOnly: true, - }); } /** @@ -3374,18 +3394,23 @@ export class AssetsController extends BaseController< * @param chains - Array of chain IDs to subscribe for. * @param options - Optional subscription overrides. * @param options.subscriptionKey - Custom subscription key (default: `ds:`). - * @param options.customAssetsOnly - When true, only poll customAssets for these chains. + * @param options.customAssets - Pinned assets this source claimed + * (`claimCustomAssets`), forwarded on the poll request. */ #subscribeDataSource( source: AbstractDataSource, accounts: InternalAccount[], chains: ChainId[], - options: { subscriptionKey?: string; customAssetsOnly?: boolean } = {}, + options: { + subscriptionKey?: string; + customAssets?: Caip19AssetId[]; + } = {}, ): void { const sourceId = source.getName(); const subscriptionKey = options.subscriptionKey ?? `ds:${sourceId}`; const existingSubscription = this.#activeSubscriptions.get(subscriptionKey); const isUpdate = existingSubscription !== undefined; + const customAssets = options.customAssets ?? []; log('Subscribe to data source', { sourceId, @@ -3393,17 +3418,19 @@ export class AssetsController extends BaseController< isUpdate, accountCount: accounts.length, chainCount: chains.length, - customAssetsOnly: options.customAssetsOnly === true, + customAssetCount: customAssets.length, }); + // Globally hidden assets, forwarded as `excludeAssetIds`. + const hiddenAssets = this.#getHiddenAssetIds(); + const subscribeReq: SubscriptionRequest = { request: this.#buildDataRequest(accounts, chains, { assetTypes: ['fungible'], dataTypes: ['balance'], updateInterval: this.#defaultUpdateInterval, - ...(options.customAssetsOnly === true - ? { customAssetsOnly: true } - : {}), + customAssets: customAssets.length > 0 ? customAssets : undefined, + excludeAssetIds: hiddenAssets.length > 0 ? hiddenAssets : undefined, }), subscriptionId: subscriptionKey, isUpdate, @@ -3853,20 +3880,14 @@ export class AssetsController extends BaseController< ), }; - // 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 === 'AccountActivityDataSource'; - - const enrichmentSources: AssetsDataSource[] = [ - ...(shouldGraduateCustomAssets - ? [this.#customAssetGraduationMiddleware] - : []), - this.#detectionMiddleware, - ]; + // Recover errored chains and unresolved pins on RPC before enrichment, + // mirroring the force-update pipeline. In RPC-only mode the poll + // already uses RPC, so there is nothing to fall back to. + const enrichmentSources: AssetsDataSource[] = []; + if (this.#isBasicFunctionality()) { + enrichmentSources.push(this.#rpcFallbackMiddleware); + } + enrichmentSources.push(this.#detectionMiddleware); if (this.#isBasicFunctionality()) { enrichmentSources.push( createParallelMiddleware([ diff --git a/packages/assets-controller/src/data-sources/AbstractDataSource.test.ts b/packages/assets-controller/src/data-sources/AbstractDataSource.test.ts index 98c5262e8e6..b4e11723bf6 100644 --- a/packages/assets-controller/src/data-sources/AbstractDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AbstractDataSource.test.ts @@ -1,4 +1,4 @@ -import type { ChainId } from '../types.js'; +import type { Caip19AssetId, ChainId } from '../types.js'; import type { ActiveSubscription, DataSourceState, @@ -113,6 +113,43 @@ describe('AbstractDataSource', () => { expect(dataSource.getSubscriptions().size).toBe(0); }); + describe('claimCustomAssets', () => { + it('claims assets on assigned chains by default and skips others', () => { + const { dataSource } = setupDataSource(); + + const assignedChainAsset = + `${CHAIN_MAINNET}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + const unassignedChainAsset = + `${CHAIN_POLYGON}/erc20:0x2222222222222222222222222222222222222222` as Caip19AssetId; + + expect( + dataSource.claimCustomAssets( + [assignedChainAsset, unassignedChainAsset], + [CHAIN_MAINNET], + ), + ).toStrictEqual([assignedChainAsset]); + }); + + it('skips malformed asset IDs and claims nothing without assigned chains', () => { + const { dataSource } = setupDataSource(); + + expect( + dataSource.claimCustomAssets( + ['not-a-caip-asset' as Caip19AssetId], + [CHAIN_MAINNET], + ), + ).toStrictEqual([]); + expect( + dataSource.claimCustomAssets( + [ + `${CHAIN_MAINNET}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId, + ], + [], + ), + ).toStrictEqual([]); + }); + }); + it.each([ { chains: [], expected: [] }, { chains: [CHAIN_MAINNET], expected: [CHAIN_MAINNET] }, diff --git a/packages/assets-controller/src/data-sources/AbstractDataSource.ts b/packages/assets-controller/src/data-sources/AbstractDataSource.ts index a669700e101..7de67a7301d 100644 --- a/packages/assets-controller/src/data-sources/AbstractDataSource.ts +++ b/packages/assets-controller/src/data-sources/AbstractDataSource.ts @@ -1,4 +1,7 @@ +import { parseCaipAssetType } from '@metamask/utils'; + import type { + Caip19AssetId, ChainId, DataRequest, DataResponse, @@ -111,6 +114,31 @@ export abstract class AbstractDataSource< return this.state.activeChains; } + /** + * Claim the pinned assets this source commits to serving. Called during the + * subscription handoff; claimed assets are not offered to lower-priority + * sources. Assets a source cannot resolve at fetch time are released via + * `DataResponse.unprocessedCustomAssets`. Default: claim assets on this + * source's assigned chains. + * + * @param customAssets - Candidate CAIP-19 asset IDs still unclaimed. + * @param assignedChains - Chains assigned to this source in the handoff. + * @returns The claimed subset of `customAssets`. + */ + claimCustomAssets( + customAssets: Caip19AssetId[], + assignedChains: ChainId[], + ): Caip19AssetId[] { + const assigned = new Set(assignedChains); + return customAssets.filter((assetId) => { + try { + return assigned.has(parseCaipAssetType(assetId).chainId); + } catch { + return false; + } + }); + } + /** * Subscribe to updates for the given request. */ diff --git a/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts index 13b701d80f5..8bf2cf4b11c 100644 --- a/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountActivityDataSource.ts @@ -245,6 +245,9 @@ export class AccountActivityDataSource extends AbstractDataSource< // SUBSCRIBE / UNSUBSCRIBE // ============================================================================ + // Pins on active chains are claimed via the inherited `claimCustomAssets`: + // the activity stream pushes all activity for the address, pins included. + /** * AADS is event-driven and chain-agnostic: it never participates in the * controller's subscribe/unsubscribe handoff. Incoming `balanceUpdated` diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index 0bec4e86f9e..849cdac4172 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -71,6 +71,7 @@ function createMockApiClient( balances: V5BalanceItem[] = [], unprocessedNetworks: string[] = [], v6Accounts: V6AccountBalancesEntry[] = [], + unprocessedIncludeAssetIds: string[] = [], ): MockApiClient { return { accounts: { @@ -85,7 +86,7 @@ function createMockApiClient( fetchV6MultiAccountBalances: jest.fn().mockResolvedValue({ accounts: v6Accounts, unprocessedNetworks, - unprocessedIncludeAssetIds: [], + unprocessedIncludeAssetIds, }), }, }; @@ -146,6 +147,7 @@ async function setupController( supportedChains?: number[]; balances?: V5BalanceItem[]; unprocessedNetworks?: string[]; + unprocessedIncludeAssetIds?: string[]; fetchTimeoutMs?: number; v6Accounts?: V6AccountBalancesEntry[]; remoteFeatureFlags?: Record; @@ -155,6 +157,7 @@ async function setupController( supportedChains = [1, 137], balances = [], unprocessedNetworks = [], + unprocessedIncludeAssetIds = [], fetchTimeoutMs, v6Accounts = [], remoteFeatureFlags = {}, @@ -198,6 +201,7 @@ async function setupController( balances, unprocessedNetworks, v6Accounts, + unprocessedIncludeAssetIds, ); const controller = new AccountsApiDataSource({ @@ -593,6 +597,48 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); + describe('claimCustomAssets', () => { + const assignedChainAsset = + 'eip155:1/erc20:0x1111111111111111111111111111111111111111'; + const unassignedChainAsset = + 'eip155:137/erc20:0x2222222222222222222222222222222222222222'; + const nonEvmAsset = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFW'; + + it('claims EVM assets on assigned chains when the v6 flag is enabled', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + expect( + controller.claimCustomAssets( + [ + assignedChainAsset, + unassignedChainAsset, + nonEvmAsset, + 'not-a-caip-asset', + ] as Caip19AssetId[], + ['eip155:1' as ChainId], + ), + ).toStrictEqual([assignedChainAsset]); + + controller.destroy(); + }); + + it('claims nothing when the v6 flag is disabled (v5 has no includeAssetIds support)', async () => { + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: false } }, + }); + + expect( + controller.claimCustomAssets([assignedChainAsset] as Caip19AssetId[], [ + 'eip155:1' as ChainId, + ]), + ).toStrictEqual([]); + + controller.destroy(); + }); + }); + describe('assetsAccountsApiV6 feature flag', () => { it('uses the v5 endpoint by default', async () => { const { controller, apiClient } = await setupController(); @@ -786,7 +832,7 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('does not pass includeAssetIds to v6 even when custom assets are present', async () => { + it('passes EVM custom assets on requested chains to v6 as includeAssetIds', async () => { const { controller, apiClient } = await setupController({ remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }); @@ -802,7 +848,252 @@ describe('AccountsApiDataSource', () => { apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledWith( [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [customToken] }, + undefined, + ); + + controller.destroy(); + }); + + it('omits custom assets that are not on a requested chain from includeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + // Custom asset on Polygon while only Mainnet is being fetched. + const polygonToken = + 'eip155:137/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ + chainIds: [CHAIN_MAINNET], + customAssets: [polygonToken], + }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + undefined, + undefined, + ); + + controller.destroy(); + }); + + it('surfaces unprocessed include asset ids on the asset axis (unprocessedCustomAssets) without flagging the chain as errored', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + unprocessedIncludeAssetIds: [customToken], + }); + + const response = await controller.fetch( + createDataRequest({ customAssets: [customToken] }), + ); + + // The chain itself succeeded — only the specific pinned asset is + // outstanding, so it goes on the asset axis, not `errors`. + expect(response.errors?.[CHAIN_MAINNET]).toBeUndefined(); + expect(response.unprocessedCustomAssets).toStrictEqual([customToken]); + + controller.destroy(); + }); + + it('omits unparseable unprocessed include asset ids from unprocessedCustomAssets', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + unprocessedIncludeAssetIds: [ + customToken, + 'not-a-caip-asset' as Caip19AssetId, + ], + }); + + const response = await controller.fetch( + createDataRequest({ customAssets: [customToken] }), + ); + + expect(response.unprocessedCustomAssets).toStrictEqual([customToken]); + + controller.destroy(); + }); + + it('skips non-EVM and malformed custom assets when building includeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + const solanaToken = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; + const malformed = 'not-a-caip-asset' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ customAssets: [solanaToken, malformed] }), + ); + + // No EVM custom asset on a requested chain -> includeAssetIds omitted. + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + undefined, undefined, + ); + + controller.destroy(); + }); + + it('ignores malformed unprocessed include asset ids', async () => { + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const { controller } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + unprocessedIncludeAssetIds: ['not-a-caip-asset'], + }); + + const response = await controller.fetch( + createDataRequest({ customAssets: [customToken] }), + ); + + // The malformed unprocessed id cannot be parsed, so it is dropped from + // both axes (no error, no asset-axis entry). + expect(response.errors).toBeUndefined(); + expect(response.unprocessedCustomAssets).toBeUndefined(); + + controller.destroy(); + }); + + it('passes EVM hidden assets on requested chains to v6 as excludeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + const hiddenToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ excludeAssetIds: [hiddenToken] }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { excludeAssetIds: [hiddenToken] }, + undefined, + ); + + controller.destroy(); + }); + + it('omits hidden assets that are not on a requested chain from excludeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + // Hidden asset on Polygon while only Mainnet is being fetched. + const polygonToken = + 'eip155:137/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ + chainIds: [CHAIN_MAINNET], + excludeAssetIds: [polygonToken], + }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + undefined, + undefined, + ); + + controller.destroy(); + }); + + it('skips non-EVM and malformed hidden assets when building excludeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + const solanaToken = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; + const malformed = 'not-a-caip-asset' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ excludeAssetIds: [solanaToken, malformed] }), + ); + + // No EVM hidden asset on a requested chain -> excludeAssetIds omitted. + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + undefined, + undefined, + ); + + controller.destroy(); + }); + + it('lets a pinned asset win when it also appears in the hidden list', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + const token = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ + customAssets: [token], + excludeAssetIds: [token], + }), + ); + + // The asset is pinned, so it is included and never excluded. + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [token] }, + undefined, + ); + + controller.destroy(); + }); + + it('sends both includeAssetIds and excludeAssetIds when pins and hidden assets differ', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, + }); + + const pinned = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const hidden = + 'eip155:1/erc20:0xdAC17F958D2ee523a2206206994597C13D831ec7' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ + customAssets: [pinned], + excludeAssetIds: [hidden], + }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [pinned], excludeAssetIds: [hidden] }, undefined, ); diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index 95339c45fc5..647614ed37c 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -10,6 +10,7 @@ import type { import { isCaipChainId, KnownCaipNamespace, + parseCaipAssetType, toCaipChainId, } from '@metamask/utils'; @@ -430,11 +431,32 @@ export class AccountsApiDataSource extends AbstractDataSource< ? { staleTime: 100, gcTime: 100 } : undefined; + // User-pinned assets on the fetched chains, sent to v6 as + // `includeAssetIds` so the backend returns them even at zero balance + // (APIPLAT-2499). + const includeAssetIds = this.#getIncludeAssetIds(request, chainsToFetch); + + // User-hidden assets on the fetched chains, sent to v6 as + // `excludeAssetIds`. A pin wins over a hide, so overlap with + // `includeAssetIds` is removed. + const excludeAssetIds = this.#getExcludeAssetIds( + request, + chainsToFetch, + includeAssetIds, + ); + // Feature-flagged: v6 endpoint with a fallback to legacy v5. The flag is // read here (not cached) so a runtime toggle can revert v6 -> v5. - const { unprocessedNetworks, assetsBalance } = this.#isBalanceV6Enabled() - ? await this.#fetchV6Balances(accountIds, fetchOptions, request) - : await this.#fetchV5Balances(accountIds, fetchOptions, request); + const { unprocessedNetworks, unprocessedIncludeAssetIds, assetsBalance } = + this.#isBalanceV6Enabled() + ? await this.#fetchV6Balances( + accountIds, + fetchOptions, + request, + includeAssetIds, + excludeAssetIds, + ) + : await this.#fetchV5Balances(accountIds, fetchOptions, request); // Handle unprocessed networks - these will be passed to next middleware if (unprocessedNetworks.length > 0) { @@ -448,6 +470,26 @@ export class AccountsApiDataSource extends AbstractDataSource< } } + // Pins the backend could not resolve go on `unprocessedCustomAssets` — + // not `errors`, the chain succeeded — so the RPC fallback fetches just + // these assets. + const validUnprocessedAssetIds = unprocessedIncludeAssetIds.filter( + (assetId) => { + try { + parseCaipAssetType(assetId as Caip19AssetId); + return true; + } catch { + return false; + } + }, + ); + if (validUnprocessedAssetIds.length > 0) { + response.unprocessedCustomAssets = [ + ...(response.unprocessedCustomAssets ?? []), + ...(validUnprocessedAssetIds as Caip19AssetId[]), + ]; + } + response.assetsBalance = assetsBalance; response.updateMode = 'merge'; } catch (error) { @@ -477,6 +519,87 @@ export class AccountsApiDataSource extends AbstractDataSource< return response; } + /** + * Collect the pinned EVM assets on the fetched chains to send to the v6 + * endpoint as `includeAssetIds`; malformed IDs are skipped. + * + * @param request - The data request (carries `customAssets`). + * @param chainsToFetch - Chains being requested this fetch. + * @returns Deduplicated asset IDs, or `undefined` when none. + */ + #getIncludeAssetIds( + request: DataRequest, + chainsToFetch: ChainId[], + ): Caip19AssetId[] | undefined { + if (!request.customAssets || request.customAssets.length === 0) { + return undefined; + } + + const chainsToFetchSet = new Set(chainsToFetch); + const includeAssetIds = new Set(); + + for (const assetId of request.customAssets) { + let chainId: ChainId; + try { + chainId = parseCaipAssetType(assetId).chainId; + } catch { + continue; + } + if ( + chainId.startsWith(`${KnownCaipNamespace.Eip155}:`) && + chainsToFetchSet.has(chainId) + ) { + includeAssetIds.add(assetId); + } + } + + return includeAssetIds.size > 0 ? [...includeAssetIds] : undefined; + } + + /** + * Collect the hidden EVM assets on the fetched chains to send to the v6 + * endpoint as `excludeAssetIds`; malformed IDs are skipped and pinned + * assets are left out (a pin wins). + * + * @param request - The data request (carries `excludeAssetIds`). + * @param chainsToFetch - Chains being requested this fetch. + * @param includeAssetIds - Pinned asset IDs that must not be excluded. + * @returns Deduplicated asset IDs, or `undefined` when none. + */ + #getExcludeAssetIds( + request: DataRequest, + chainsToFetch: ChainId[], + includeAssetIds: Caip19AssetId[] | undefined, + ): Caip19AssetId[] | undefined { + if (!request.excludeAssetIds || request.excludeAssetIds.length === 0) { + return undefined; + } + + const chainsToFetchSet = new Set(chainsToFetch); + const includeSet = new Set(includeAssetIds ?? []); + const excludeAssetIds = new Set(); + + for (const assetId of request.excludeAssetIds) { + if (includeSet.has(assetId)) { + continue; + } + let chainId: ChainId; + try { + chainId = parseCaipAssetType(assetId).chainId; + } catch { + continue; + } + if ( + chainId.startsWith(`${KnownCaipNamespace.Eip155}:`) && + chainsToFetchSet.has(chainId) + ) { + excludeAssetIds.add(assetId); + } + } + + return excludeAssetIds.size > 0 ? [...excludeAssetIds] : undefined; + } + /** * Fetch balances from the legacy v5 endpoint and process them. * @@ -491,6 +614,7 @@ export class AccountsApiDataSource extends AbstractDataSource< request: DataRequest, ): Promise<{ unprocessedNetworks: string[]; + unprocessedIncludeAssetIds: string[]; assetsBalance: Record>; }> { const apiResponse = await fetchWithTimeout( @@ -510,6 +634,8 @@ export class AccountsApiDataSource extends AbstractDataSource< return { unprocessedNetworks: apiResponse.unprocessedNetworks, + // v5 has no `includeAssetIds` support. + unprocessedIncludeAssetIds: [], assetsBalance, }; } @@ -520,21 +646,35 @@ export class AccountsApiDataSource extends AbstractDataSource< * @param accountIds - CAIP-10 account IDs to fetch balances for. * @param fetchOptions - Cache/fetch options (e.g. force update settings). * @param request - The original data request containing accounts to map. - * @returns Unprocessed networks and processed asset balances by account. + * @param includeAssetIds - Pinned asset IDs the backend must always return. + * @param excludeAssetIds - Hidden asset IDs the backend must drop. + * @returns Unprocessed networks, unprocessed pinned assets, and processed + * asset balances by account. */ async #fetchV6Balances( accountIds: string[], fetchOptions: { staleTime: number; gcTime: number } | undefined, request: DataRequest, + includeAssetIds: Caip19AssetId[] | undefined, + excludeAssetIds: Caip19AssetId[] | undefined, ): Promise<{ unprocessedNetworks: string[]; + unprocessedIncludeAssetIds: string[]; assetsBalance: Record>; }> { + const params = + includeAssetIds || excludeAssetIds + ? { + ...(includeAssetIds && { includeAssetIds }), + ...(excludeAssetIds && { excludeAssetIds }), + } + : undefined; + const apiResponse = await fetchWithTimeout( () => this.#apiClient.accounts.fetchV6MultiAccountBalances( accountIds, - undefined, + params, fetchOptions, ), this.#fetchTimeoutMs, @@ -547,6 +687,7 @@ export class AccountsApiDataSource extends AbstractDataSource< return { unprocessedNetworks: apiResponse.unprocessedNetworks, + unprocessedIncludeAssetIds: apiResponse.unprocessedIncludeAssetIds, assetsBalance, }; } @@ -729,6 +870,17 @@ export class AccountsApiDataSource extends AbstractDataSource< } } + // Forward the asset-axis signal so the RPC fallback recovers these pins. + if ( + response.unprocessedCustomAssets && + response.unprocessedCustomAssets.length > 0 + ) { + context.response.unprocessedCustomAssets = [ + ...(context.response.unprocessedCustomAssets ?? []), + ...response.unprocessedCustomAssets, + ]; + } + // Determine successfully handled chains (exclude unprocessed/error chains) const unprocessedChains = new Set(Object.keys(response.errors ?? {})); successfullyHandledChains = request.chainIds.filter( @@ -774,6 +926,36 @@ export class AccountsApiDataSource extends AbstractDataSource< // SUBSCRIBE // ============================================================================ + /** + * Claim EVM pins on assigned chains (sent to v6 as `includeAssetIds`). + * v5 has no `includeAssetIds`, so with the v6 flag off nothing is claimed + * and pins fall through to RPC. + * + * @param customAssets - Candidate CAIP-19 asset IDs still unclaimed. + * @param assignedChains - Chains assigned to this source in the handoff. + * @returns The claimed subset of `customAssets`. + */ + claimCustomAssets( + customAssets: Caip19AssetId[], + assignedChains: ChainId[], + ): Caip19AssetId[] { + if (!this.#isBalanceV6Enabled()) { + return []; + } + const assigned = new Set(assignedChains); + return customAssets.filter((assetId) => { + try { + const parsed = parseCaipAssetType(assetId); + return ( + parsed.chain.namespace === KnownCaipNamespace.Eip155 && + assigned.has(parsed.chainId) + ); + } catch { + return false; + } + }); + } + async subscribe(subscriptionRequest: SubscriptionRequest): Promise { const { request, subscriptionId, isUpdate } = subscriptionRequest; diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts index 80d2e9edde3..bd2bc541650 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.test.ts @@ -832,28 +832,38 @@ describe('RpcDataSource', () => { .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') .mockResolvedValue(createBalanceFetchResult()); - await withController(async ({ controller }) => { - const request = createDataRequest({ - customAssets: [customAssetId], - }); - await controller.fetch(request); - - expect(fetchSpy).toHaveBeenCalledWith( - MOCK_CHAIN_ID_HEX, - MOCK_ACCOUNT_ID, - MOCK_ADDRESS, - [ - { - assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, - address: '0x0000000000000000000000000000000000000000', - }, - expect.objectContaining({ - assetId: customAssetId, - address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [customAssetId] }, }), - ], - ); - }); + }, + }, + async ({ controller }) => { + const request = createDataRequest({ + customAssets: [customAssetId], + }); + await controller.fetch(request); + + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + expect.objectContaining({ + assetId: customAssetId, + address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + }), + ], + ); + }, + ); fetchSpy.mockRestore(); }); @@ -868,28 +878,83 @@ describe('RpcDataSource', () => { .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') .mockResolvedValue(createBalanceFetchResult()); - await withController(async ({ controller }) => { - const request = createDataRequest({ - customAssets: [matchingAsset, otherChainAsset], - }); - await controller.fetch(request); + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { + [MOCK_ACCOUNT_ID]: [matchingAsset, otherChainAsset], + }, + }), + }, + }, + async ({ controller }) => { + const request = createDataRequest({ + customAssets: [matchingAsset, otherChainAsset], + }); + await controller.fetch(request); - expect(fetchSpy).toHaveBeenCalledWith( - MOCK_CHAIN_ID_HEX, - MOCK_ACCOUNT_ID, - MOCK_ADDRESS, - [ - { - assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, - address: '0x0000000000000000000000000000000000000000', - }, - expect.objectContaining({ - assetId: matchingAsset, - address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + expect.objectContaining({ + assetId: matchingAsset, + address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', + }), + ], + ); + }, + ); + + fetchSpy.mockRestore(); + }); + + it('does not fetch custom assets the fetching account has not pinned', async () => { + const customAssetId = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + const fetchSpy = jest + .spyOn(BalanceFetcher.prototype, 'fetchBalancesForAssets') + .mockResolvedValue(createBalanceFetchResult()); + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { 'other-account-id': [customAssetId] }, }), - ], - ); - }); + }, + }, + async ({ controller }) => { + const request = createDataRequest({ + customAssets: [customAssetId], + }); + await controller.fetch(request); + + // The request's customAssets list is flat; ownership comes from + // state. MOCK_ACCOUNT_ID did not pin the asset, so its fetch + // carries only the native entry. + expect(fetchSpy).toHaveBeenCalledWith( + MOCK_CHAIN_ID_HEX, + MOCK_ACCOUNT_ID, + MOCK_ADDRESS, + [ + { + assetId: `${MOCK_CHAIN_ID_CAIP}/slip44:60`, + address: '0x0000000000000000000000000000000000000000', + }, + ], + ); + }, + ); fetchSpy.mockRestore(); }); @@ -1358,6 +1423,209 @@ describe('RpcDataSource', () => { expect(detectionStopSpy).toHaveBeenCalled(); }); }); + + it('starts an asset-scoped poll (explicit assetIds) for pinned assets on chains not assigned to RPC', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const detectionStartSpy = jest.spyOn( + TokenDetector.prototype, + 'startPolling', + ); + const customAssetId = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [customAssetId] }, + }), + }, + }, + async ({ controller }) => { + // The chain axis assigned nothing to RPC (chainIds is empty — e.g. + // the websocket claimed the chain), but RPC claimed the pinned asset. + await controller.subscribe({ + request: createDataRequest({ + chainIds: [], + customAssets: [customAssetId], + }), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + + expect(balanceStartSpy).toHaveBeenCalledWith({ + chainId: MOCK_CHAIN_ID_HEX, + accountId: MOCK_ACCOUNT_ID, + accountAddress: MOCK_ADDRESS, + assetIds: [customAssetId], + }); + // No regular chain coverage — detection must not run. + expect(detectionStartSpy).not.toHaveBeenCalled(); + await controller.unsubscribe('test-sub'); + }, + ); + }); + + it('does not start an asset-scoped poll for accounts without pinned assets on the supplemental chain', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const customAssetId = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { 'other-account-id': [customAssetId] }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest({ + chainIds: [], + customAssets: [customAssetId], + }), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + + expect(balanceStartSpy).not.toHaveBeenCalled(); + await controller.unsubscribe('test-sub'); + }, + ); + }); + + it('does not start an asset-scoped poll for chains already covered by regular polling', async () => { + const balanceStartSpy = jest.spyOn( + BalanceFetcher.prototype, + 'startPolling', + ); + const customAssetId = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [customAssetId] }, + }), + }, + }, + async ({ controller }) => { + // Chain assigned to RPC: the regular poll already includes + // state.customAssets, so no supplemental poll must start. + await controller.subscribe({ + request: createDataRequest({ + customAssets: [customAssetId], + }), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + + expect(balanceStartSpy).toHaveBeenCalledTimes(1); + expect(balanceStartSpy).toHaveBeenCalledWith({ + chainId: MOCK_CHAIN_ID_HEX, + accountId: MOCK_ACCOUNT_ID, + accountAddress: MOCK_ADDRESS, + }); + await controller.unsubscribe('test-sub'); + }, + ); + }); + + it('unsubscribe stops asset-scoped polling', async () => { + const balanceStopSpy = jest.spyOn( + BalanceFetcher.prototype, + 'stopPollingByPollingToken', + ); + const customAssetId = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + + await withController( + { + actionHandlerOverrides: { + 'AssetsController:getState': () => ({ + ...getDefaultAssetsControllerState(), + customAssets: { [MOCK_ACCOUNT_ID]: [customAssetId] }, + }), + }, + }, + async ({ controller }) => { + await controller.subscribe({ + request: createDataRequest({ + chainIds: [], + customAssets: [customAssetId], + }), + subscriptionId: 'test-sub', + isUpdate: false, + onAssetsUpdate: jest.fn(), + }); + await controller.unsubscribe('test-sub'); + expect(balanceStopSpy).toHaveBeenCalled(); + }, + ); + }); + }); + + describe('claimCustomAssets', () => { + const availableChainAsset = + `${MOCK_CHAIN_ID_CAIP}/erc20:0x1111111111111111111111111111111111111111` as Caip19AssetId; + const unavailableChainAsset = + 'eip155:999/erc20:0x2222222222222222222222222222222222222222' as Caip19AssetId; + const nonEvmAsset = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFW' as Caip19AssetId; + + it('claims EVM assets on active chains even when the chain was not assigned to RPC', async () => { + await withController(async ({ controller }) => { + expect( + controller.claimCustomAssets([availableChainAsset], []), + ).toStrictEqual([availableChainAsset]); + }); + }); + + it('does not claim assets on chains RPC cannot serve, non-EVM assets, or malformed IDs', async () => { + await withController(async ({ controller }) => { + expect( + controller.claimCustomAssets( + [ + unavailableChainAsset, + nonEvmAsset, + 'not-a-caip-asset' as Caip19AssetId, + ], + [], + ), + ).toStrictEqual([]); + }); + }); + + it('falls back to assigned chains when network state has not been applied yet', async () => { + const networkState = createMockNetworkState(NetworkStatus.Degraded); + await withController({ networkState }, async ({ controller }) => { + // eslint-disable-next-line n/no-sync -- testing sync API used by AssetsController + expect(controller.getActiveChainsSync()).toStrictEqual([]); + expect( + controller.claimCustomAssets( + [availableChainAsset], + [MOCK_CHAIN_ID_CAIP], + ), + ).toStrictEqual([availableChainAsset]); + expect( + controller.claimCustomAssets([availableChainAsset], []), + ).toStrictEqual([]); + }); + }); }); describe('unsubscribe', () => { diff --git a/packages/assets-controller/src/data-sources/RpcDataSource.ts b/packages/assets-controller/src/data-sources/RpcDataSource.ts index 8b3a7795db2..f818abc5537 100644 --- a/packages/assets-controller/src/data-sources/RpcDataSource.ts +++ b/packages/assets-controller/src/data-sources/RpcDataSource.ts @@ -15,11 +15,11 @@ import type { import { isStrictHexString, isCaipChainId, + KnownCaipNamespace, parseCaipAssetType, parseCaipChainId, hexToNumber, toCaipChainId, - KnownCaipNamespace, } from '@metamask/utils'; import type { Hex } from '@metamask/utils'; import BigNumberJS from 'bignumber.js'; @@ -1015,6 +1015,12 @@ export class RpcDataSource extends AbstractDataSource< const assetsInfo: Record = {}; const failedChains: ChainId[] = []; + // request.customAssets is flat — resolve pin ownership from state so + // each account only fetches its own pins. + const customAssetsByAccount = request.customAssets + ? this.#getCustomAssetsByAccount() + : {}; + // Fetch balances for each account and its supported chains (pre-computed in request) for (const { account, @@ -1028,6 +1034,9 @@ export class RpcDataSource extends AbstractDataSource< } const { address, id: accountId } = account; + const pinnedByAccount = new Set( + (customAssetsByAccount[accountId] ?? []).map(normalizeAssetId), + ); for (const chainId of chainsForAccount) { const hexChainId = caipChainIdToHex(chainId); @@ -1047,13 +1056,14 @@ export class RpcDataSource extends AbstractDataSource< try { const parsed = parseCaipAssetType(assetId); const assetChainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; + const normalizedId = normalizeAssetId(assetId); if ( assetChainId === chainId && + pinnedByAccount.has(normalizedId) && this.#getAssetType(assetId) === 'erc20' ) { const tokenAddress = parsed.assetReference.toLowerCase() as Address; - const normalizedId = normalizeAssetId(assetId); const decimals = existingMetadata[normalizedId]?.decimals; assetsToFetch.push({ @@ -1360,6 +1370,37 @@ export class RpcDataSource extends AbstractDataSource< }; } + /** + * RPC is the terminal claimer on the asset axis: it claims every EVM pin it + * has a provider for, even on chains claimed by higher-priority sources. + * Pins outside the regular RPC assignment get an asset-scoped poll (see + * `subscribe`). + * + * @param customAssets - Candidate CAIP-19 asset IDs still unclaimed. + * @param assignedChains - Chains assigned to RPC; availability fallback + * before network state is applied. + * @returns The claimed subset of `customAssets`. + */ + claimCustomAssets( + customAssets: Caip19AssetId[], + assignedChains: ChainId[], + ): Caip19AssetId[] { + const available = new Set( + this.#activeChains.length > 0 ? this.#activeChains : assignedChains, + ); + return customAssets.filter((assetId) => { + try { + const parsed = parseCaipAssetType(assetId); + return ( + parsed.chain.namespace === KnownCaipNamespace.Eip155 && + available.has(parsed.chainId) + ); + } catch { + return false; + } + }); + } + /** * Subscribe to updates for the given request. * Starts polling through BalanceFetcher and TokenDetector. @@ -1383,15 +1424,23 @@ export class RpcDataSource extends AbstractDataSource< ) : request.chainIds; + // Pins claimed on chains outside the regular RPC assignment get an + // asset-scoped poll below. + const supplementalChains = this.#getSupplementalCustomAssetChains( + request, + chainsToSubscribe, + ); + log('Subscribe requested', { subscriptionId, isUpdate, accounts: request.accountsWithSupportedChains.map((a) => a.account.id), chainsToSubscribe, + supplementalChains, activeChainsFallback: this.#activeChains.length === 0, }); - if (chainsToSubscribe.length === 0) { + if (chainsToSubscribe.length === 0 && supplementalChains.length === 0) { log('No active chains to subscribe'); return; } @@ -1436,20 +1485,11 @@ export class RpcDataSource extends AbstractDataSource< chainId: hexChainId, accountId, accountAddress: address as Address, - ...(request.customAssetsOnly === true - ? { customAssetsOnly: true } - : {}), }; const balanceToken = this.#balanceFetcher.startPolling(balanceInput); balancePollingTokens.push(balanceToken); - // Token detection is only relevant for "regular" subscriptions — - // a customAssetsOnly subscription should never run detection. - if ( - request.customAssetsOnly !== true && - this.#tokenDetectionEnabled() && - this.#useExternalService() - ) { + if (this.#tokenDetectionEnabled() && this.#useExternalService()) { const detectionInput: DetectionPollingInput = { chainId: hexChainId, accountId, @@ -1462,6 +1502,52 @@ export class RpcDataSource extends AbstractDataSource< } } + // Asset-scoped polls on chains another source claimed: poll ONLY the + // claimed pins to avoid double-polling tracked balances. Pin changes + // re-run the subscription pass, which rebuilds these polls. + if (supplementalChains.length > 0) { + const supplemental = new Set(supplementalChains); + const claimedAssetsByChain = new Map(); + for (const assetId of request.customAssets ?? []) { + try { + const { chainId } = parseCaipAssetType(assetId); + if (supplemental.has(chainId)) { + const chainAssets = claimedAssetsByChain.get(chainId) ?? []; + chainAssets.push(assetId); + claimedAssetsByChain.set(chainId, chainAssets); + } + } catch { + // Skip unparseable asset IDs + } + } + // request.customAssets is flat; ownership comes from controller state. + const customAssetsByAccount = this.#getCustomAssetsByAccount(); + for (const { account } of request.accountsWithSupportedChains) { + const pinned = new Set(customAssetsByAccount[account.id] ?? []); + if (pinned.size === 0) { + continue; + } + for (const [chainId, chainAssets] of claimedAssetsByChain) { + // Sorted so the polling input (the dedupe key) is deterministic. + const assetIds = chainAssets + .filter((assetId) => pinned.has(assetId)) + .sort(); + if (assetIds.length === 0) { + continue; + } + const balanceInput: BalancePollingInput = { + chainId: caipChainIdToHex(chainId), + accountId: account.id, + accountAddress: account.address as Address, + assetIds, + }; + balancePollingTokens.push( + this.#balanceFetcher.startPolling(balanceInput), + ); + } + } + } + // Store subscription data const accounts = request.accountsWithSupportedChains.map( (entry) => entry.account, @@ -1505,6 +1591,63 @@ export class RpcDataSource extends AbstractDataSource< } } + /** + * Chains needing a supplemental asset-scoped poll: chains of pins not + * covered by the regular RPC polling. Only EVM chains RPC can serve; + * malformed IDs are skipped. + * + * @param request - The subscription's data request (carries `customAssets`). + * @param chainsToSubscribe - Chains covered by the regular polling loop. + * @returns Chains requiring an asset-scoped poll. + */ + #getSupplementalCustomAssetChains( + request: DataRequest, + chainsToSubscribe: ChainId[], + ): ChainId[] { + if (!request.customAssets || request.customAssets.length === 0) { + return []; + } + + const covered = new Set(chainsToSubscribe); + const chains = new Set(); + + for (const assetId of request.customAssets) { + let parsed: ReturnType; + try { + parsed = parseCaipAssetType(assetId); + } catch { + continue; + } + const { chainId } = parsed; + if ( + parsed.chain.namespace === KnownCaipNamespace.Eip155 && + !covered.has(chainId) && + (this.#activeChains.length === 0 || + this.#activeChains.includes(chainId)) + ) { + chains.add(chainId); + } + } + + return [...chains]; + } + + /** + * Get per-account pins from AssetsController state — the request's flat + * `customAssets` list carries no account association. + * + * @returns Record of account ID to pinned CAIP-19 asset IDs. + */ + #getCustomAssetsByAccount(): Record { + try { + const state = this.#messenger.call('AssetsController:getState'); + return (state.customAssets ?? {}) as Record; + } catch (error) { + log('Failed to get customAssets from state', { error }); + return {}; + } + } + /** * Get existing assets metadata from AssetsController state. * Used to include metadata for ERC20 tokens when returning balance updates. diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts index 795b6f92b9b..5c02f781f7f 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.test.ts @@ -248,12 +248,10 @@ describe('BalanceFetcher', () => { ); }); - it('polls custom assets even when they have no entry in assetsBalance yet', async () => { - // Regression: custom assets must be polled by RPC because RPC is the - // sole balance fetcher for them. Previously #getAssetsToFetch only - // looked at state.assetsBalance, so a freshly added custom asset - // (no balance row yet, e.g. zero balance or first fetch failed) - // would never be polled. + it('sources the regular poll from assetsBalance only, ignoring state.customAssets', async () => { + // Pins have seeded assetsBalance rows, so the regular poll reads only + // assetsBalance; a pin without a row is NOT fetched here (asset-scoped + // polls handle those). const mockState: AssetsBalanceState = { assetsBalance: { [TEST_ACCOUNT_ID]: { @@ -283,12 +281,6 @@ describe('BalanceFetcher', () => { true, '1000000000000000000', ), - createMockBalanceResponse( - TEST_TOKEN_1.toLowerCase() as Address, - TEST_ACCOUNT, - true, - '500', - ), ]); const input: BalancePollingInput = { @@ -299,9 +291,6 @@ describe('BalanceFetcher', () => { await controller._executePoll(input); - // The multicall batch should include both the native asset and - // the custom ERC-20 token, even though the custom token has no - // entry in assetsBalance. const [, batchedRequests] = mockMulticallClient.batchBalanceOf.mock.calls[0]; const requestedTokens = ( @@ -309,17 +298,14 @@ describe('BalanceFetcher', () => { ) .map((req) => req.tokenAddress.toLowerCase()) .sort(); - expect(requestedTokens).toStrictEqual( - [ZERO_ADDRESS.toLowerCase(), TEST_TOKEN_1.toLowerCase()].sort(), - ); + expect(requestedTokens).toStrictEqual([ZERO_ADDRESS.toLowerCase()]); }, ); }); - it('in customAssetsOnly mode skips state.assetsBalance and only fetches state.customAssets', async () => { - // The supplemental subscription path: another data source covers the - // chain for regular balances, but RPC must still poll the user's - // customAssets. We must NOT also poll the regular tracked balances. + it('with explicit assetIds fetches exactly those assets and ignores tracked state', async () => { + // Asset-scoped poll path: RPC claimed the pins on a chain another + // source owns — it must NOT also poll the tracked balances. const mockState: AssetsBalanceState = { assetsBalance: { [TEST_ACCOUNT_ID]: { @@ -354,7 +340,7 @@ describe('BalanceFetcher', () => { chainId: MAINNET_CHAIN_ID, accountId: TEST_ACCOUNT_ID, accountAddress: TEST_ACCOUNT, - customAssetsOnly: true, + assetIds: [TOKEN_1_ASSET_ID], }; await controller._executePoll(input); @@ -366,8 +352,8 @@ describe('BalanceFetcher', () => { ) .map((req) => req.tokenAddress.toLowerCase()) .sort(); - // ONLY the custom token — not the native and not TOKEN_2 from - // assetsBalance. + // ONLY the pinned asset — not the native and not TOKEN_2 from + // assetsBalance, even though both are tracked in state. expect(requestedTokens).toStrictEqual([TEST_TOKEN_1.toLowerCase()]); }, ); diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts index 44d8af36fa1..9e8048a242e 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/services/BalanceFetcher.ts @@ -46,12 +46,11 @@ export type BalancePollingInput = { /** Account address */ accountAddress: Address; /** - * When true, only fetch balances for entries in `state.customAssets`, - * skipping `state.assetsBalance`. Used by the supplemental RPC - * subscription on chains that another data source is already covering - * for regular balance refreshes. + * When present, fetch exactly these assets instead of the tracked + * `assetsBalance` entries (asset-scoped RPC poll). Must be + * deterministically ordered — the polling input is the dedupe key. */ - customAssetsOnly?: boolean; + assetIds?: CaipAssetType[]; }; /** @@ -121,7 +120,7 @@ export class BalanceFetcher extends StaticIntervalPollingControllerOnly 0) { @@ -135,16 +134,15 @@ export class BalanceFetcher extends StaticIntervalPollingControllerOnly { - const assets = this.#getAssetsToFetch(chainId, accountId, customAssetsOnly); + const assets = this.#getAssetsToFetch(chainId, accountId, assetIds); return this.fetchBalancesForAssets( chainId, diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/types/state.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/types/state.ts index af114458eaf..1149e6c5eef 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/types/state.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/types/state.ts @@ -55,8 +55,8 @@ export type AssetsBalanceState = { assetsBalance: Record>; /** * User-added custom assets per account: accountId -> CAIP-19 asset IDs. - * Used to ensure RPC polling refreshes custom assets even when they have - * no entry in `assetsBalance` yet (e.g. zero balance, first fetch failed). + * Not read by the balance poll — pins have seeded `assetsBalance` rows, + * and asset-scoped polls get assets via `BalancePollingInput.assetIds`. */ customAssets?: Record; }; diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index fca7ca96476..d78109f2104 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -166,14 +166,10 @@ export type { // Middlewares export { - CustomAssetGraduationMiddleware, DetectionMiddleware, RpcFallbackMiddleware, } from './middlewares/index.js'; -export type { - CustomAssetGraduationMiddlewareOptions, - RpcFallbackMiddlewareOptions, -} from './middlewares/index.js'; +export type { RpcFallbackMiddlewareOptions } from './middlewares/index.js'; // Utilities export { diff --git a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts deleted file mode 100644 index 0b3613fe2d0..00000000000 --- a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.test.ts +++ /dev/null @@ -1,510 +0,0 @@ -import type { InternalAccount } from '@metamask/keyring-internal-api'; - -import type { - AssetsControllerStateInternal, - Caip19AssetId, - Context, - DataRequest, -} from '../types.js'; -import { CustomAssetGraduationMiddleware } from './CustomAssetGraduationMiddleware.js'; - -const MOCK_ACCOUNT_ID = 'mock-account-id'; -const OTHER_ACCOUNT_ID = 'other-account-id'; - -// Checksummed addresses — customAssets state stores normalized IDs. -const EVM_CUSTOM_ASSET = - 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; -const EVM_OTHER_ASSET = - 'eip155:137/erc20:0xdAC17F958D2ee523a2206206994597C13D831ec7' as Caip19AssetId; -const SOLANA_CUSTOM_ASSET = - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId; -const BTC_CUSTOM_ASSET = - 'bip122:000000000019d6689c085ae165831e93/slip44:0' as Caip19AssetId; - -function createMockAccount(id = MOCK_ACCOUNT_ID): InternalAccount { - return { - id, - address: '0x1234567890123456789012345678901234567890', - options: {}, - methods: [], - type: 'eip155:eoa', - scopes: ['eip155:0'], - metadata: { - name: 'Test Account', - keyring: { type: 'HD Key Tree' }, - importTime: 0, - lastSelected: 0, - }, - } as InternalAccount; -} - -function createDataRequest(overrides?: Partial): DataRequest { - const chainIds = overrides?.chainIds ?? ['eip155:1']; - const accounts = [createMockAccount()]; - return { - chainIds, - accountsWithSupportedChains: accounts.map((a) => ({ - account: a, - supportedChains: chainIds, - })), - dataTypes: ['balance'], - ...overrides, - } as DataRequest; -} - -function createAssetsState( - customAssets: Record = {}, -): AssetsControllerStateInternal { - return { - assetsInfo: {}, - assetsBalance: {}, - assetsPrice: {}, - customAssets, - assetPreferences: {}, - } as AssetsControllerStateInternal; -} - -function createContext( - overrides?: Partial, - customAssets: Record = {}, -): Context { - return { - request: createDataRequest(), - response: {}, - getAssetsState: jest.fn().mockReturnValue(createAssetsState(customAssets)), - ...overrides, - }; -} - -function setup( - customAssets: Record = {}, - selectedAccountId: string | undefined = MOCK_ACCOUNT_ID, -): { - middleware: CustomAssetGraduationMiddleware; - context: Context; - removeCustomAsset: jest.Mock; - getSelectedAccountId: jest.Mock; -} { - const removeCustomAsset = jest.fn(); - const getSelectedAccountId = jest.fn().mockReturnValue(selectedAccountId); - const middleware = new CustomAssetGraduationMiddleware({ - getSelectedAccountId, - removeCustomAsset, - }); - const context = createContext({}, customAssets); - return { middleware, context, removeCustomAsset, getSelectedAccountId }; -} - -describe('CustomAssetGraduationMiddleware', () => { - afterEach(() => { - jest.clearAllMocks(); - }); - - it('initializes with correct name', () => { - const { middleware } = setup(); - expect(middleware.name).toBe('CustomAssetGraduationMiddleware'); - expect(middleware.getName()).toBe('CustomAssetGraduationMiddleware'); - }); - - it('exposes an assetsMiddleware function', () => { - const { middleware } = setup(); - expect(typeof middleware.assetsMiddleware).toBe('function'); - }); - - it('graduates an EVM custom asset that was returned in the balance response', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(next).toHaveBeenCalledWith(context); - expect(removeCustomAsset).toHaveBeenCalledTimes(1); - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - EVM_CUSTOM_ASSET, - ); - }); - - it('graduates only the returned subset of custom assets', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET, EVM_OTHER_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).toHaveBeenCalledTimes(1); - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - EVM_CUSTOM_ASSET, - ); - }); - - it('does not graduate when AccountsAPI returns a zero balance', async () => { - // The API may include zero entries for tokens it indexes but the user - // no longer holds. Keeping them in customAssets ensures RPC keeps - // polling so a future incoming transfer is reflected immediately. - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '0' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('does not graduate when the decimal amount is zero', async () => { - // 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], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '0.0' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('graduates when AccountsAPI returns a small positive decimal balance (V5 format)', async () => { - // V5 returns amounts already divided by decimals, e.g. WETH: - // { balance: "0.283549083429656057", decimals: 18 }. The middleware - // must recognise these as positive without any further unit handling. - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '0.283549083429656057' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).toHaveBeenCalledTimes(1); - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - EVM_CUSTOM_ASSET, - ); - }); - - it('graduates only custom assets whose balance is positive', async () => { - // Mixed response: one asset has zero balance, another has a real - // balance. Only the latter should graduate. - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET, EVM_OTHER_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '0' }, - [EVM_OTHER_ASSET]: { amount: '500' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).toHaveBeenCalledTimes(1); - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - EVM_OTHER_ASSET, - ); - }); - - it('does not graduate when the balance amount is malformed', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: 'not-a-number' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('does not graduate non-EVM (Solana) custom assets', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [SOLANA_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [SOLANA_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('does not graduate non-EVM (BTC) custom assets', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [BTC_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [BTC_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('only graduates assets for the selected account', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - [OTHER_ACCOUNT_ID]: [EVM_OTHER_ASSET], - }); - context.response = { - assetsBalance: { - [OTHER_ACCOUNT_ID]: { - [EVM_OTHER_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('is a no-op when the selected account has no custom assets', async () => { - const { middleware, context, removeCustomAsset } = setup({}); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('is a no-op when the response has no balances for the selected account', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: {}, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('is a no-op when the response is empty', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = {}; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('is a no-op when there is no selected account', async () => { - const { middleware, context, removeCustomAsset, getSelectedAccountId } = - setup({ [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET] }); - getSelectedAccountId.mockReturnValue(undefined); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('does not graduate non-custom EVM assets that appear in the response', async () => { - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_OTHER_ASSET]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('does not run for non-balance data types', async () => { - const { middleware, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - const context = createContext( - { - request: createDataRequest({ dataTypes: ['metadata'] }), - response: { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }, - }, - { [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET] }, - ); - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - expect(next).toHaveBeenCalledWith(context); - }); - - it('graduates a custom asset when the response uses a non-checksummed (lowercase) address', async () => { - // 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 = - 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; - const lowercaseFromWebsocket = - 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48' as Caip19AssetId; - - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [checksummedCustomAsset], - }); - context.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [lowercaseFromWebsocket]: { amount: '1000' }, - }, - }, - }; - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).toHaveBeenCalledTimes(1); - // Removal must use the canonical (checksummed) form stored in state. - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - checksummedCustomAsset, - ); - }); - - it('does not graduate when the matching balance is added by downstream middleware (e.g. RPC fallback)', async () => { - // Regression test: the graduation middleware must inspect the response - // BEFORE calling next() so that balances merged in by later middleware - // (notably the RPC fallback, which intentionally fetches custom assets) - // do not trigger graduation. See PR description for the resilience work. - const { middleware, context, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - // No balance from upstream sources — AccountsApi did not return it. - context.response = { assetsBalance: { [MOCK_ACCOUNT_ID]: {} } }; - // The downstream middleware (RPC fallback) populates the asset balance. - const next = jest.fn().mockImplementation(async (ctx) => { - ctx.response = { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }; - return ctx; - }); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).not.toHaveBeenCalled(); - }); - - it('runs when dataTypes includes balance among others', async () => { - const { middleware, removeCustomAsset } = setup({ - [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET], - }); - const context = createContext( - { - request: createDataRequest({ dataTypes: ['balance', 'metadata'] }), - response: { - assetsBalance: { - [MOCK_ACCOUNT_ID]: { - [EVM_CUSTOM_ASSET]: { amount: '1000' }, - }, - }, - }, - }, - { [MOCK_ACCOUNT_ID]: [EVM_CUSTOM_ASSET] }, - ); - const next = jest.fn().mockImplementation((ctx) => Promise.resolve(ctx)); - - await middleware.assetsMiddleware(context, next); - - expect(removeCustomAsset).toHaveBeenCalledWith( - MOCK_ACCOUNT_ID, - EVM_CUSTOM_ASSET, - ); - }); -}); diff --git a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts b/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts deleted file mode 100644 index 17070e3cdde..00000000000 --- a/packages/assets-controller/src/middlewares/CustomAssetGraduationMiddleware.ts +++ /dev/null @@ -1,160 +0,0 @@ -import { KnownCaipNamespace } from '@metamask/utils'; - -import { projectLogger, createModuleLogger } from '../logger.js'; -import { forDataTypes } from '../types.js'; -import type { - AccountId, - AssetBalance, - Caip19AssetId, - Middleware, -} from '../types.js'; -import { normalizeAssetId } from '../utils/index.js'; - -const CONTROLLER_NAME = 'CustomAssetGraduationMiddleware'; - -const log = createModuleLogger(projectLogger, CONTROLLER_NAME); - -export type CustomAssetGraduationMiddlewareOptions = { - getSelectedAccountId: () => AccountId | undefined; - removeCustomAsset: (accountId: AccountId, assetId: Caip19AssetId) => void; -}; - -/** - * CustomAssetGraduationMiddleware removes EVM assets from `customAssets` when - * an upstream balance source (AccountsAPI / Websocket) reports a non-zero - * balance for them. Once a detector sees the asset with a real balance, it - * no longer needs to be tracked as "custom" — the regular detection flow - * will keep it fresh. - * - * Rules: - * - Only the selected account's custom assets are considered. Switching the - * selected account triggers a fresh fetch, which re-runs graduation - * against the new account's balances. - * - Only EVM (CAIP-2 namespace `eip155`) assets graduate. Non-EVM custom - * assets (Solana, BTC, Tron, etc. — served by Snap data sources) are left - * alone. - * - Only positive balances graduate. A zero balance from AccountsAPI means - * the API knows about the token but the user does not currently hold it; - * keeping it in `customAssets` ensures RPC keeps polling so a future - * incoming transfer is reflected promptly. - */ -export class CustomAssetGraduationMiddleware { - readonly name = CONTROLLER_NAME; - - readonly #getSelectedAccountId: () => AccountId | undefined; - - readonly #removeCustomAsset: ( - accountId: AccountId, - assetId: Caip19AssetId, - ) => void; - - constructor(options: CustomAssetGraduationMiddlewareOptions) { - this.#getSelectedAccountId = options.getSelectedAccountId; - this.#removeCustomAsset = options.removeCustomAsset; - } - - getName(): string { - return this.name; - } - - get assetsMiddleware(): Middleware { - return forDataTypes(['balance'], async (ctx, next) => { - // Inspect the response BEFORE calling next() so we only consider - // balances populated by upstream middleware (AccountsApi / Websocket / - // Staked). This middleware is positioned in the pipeline before the - // RPC fallback — RPC intentionally carries custom assets and must - // never trigger graduation. - const accountId = this.#getSelectedAccountId(); - if (!accountId) { - return next(ctx); - } - - const state = ctx.getAssetsState(); - const customForAccount = state.customAssets?.[accountId] ?? []; - if (customForAccount.length === 0) { - return next(ctx); - } - - const returnedBalances = ctx.response.assetsBalance?.[accountId] ?? {}; - const returnedAssetIds = Object.keys(returnedBalances) as Caip19AssetId[]; - if (returnedAssetIds.length === 0) { - return next(ctx); - } - - // customAssets state is stored with checksummed/normalized asset IDs. - // AccountsApiDataSource normalizes its response IDs, but - // 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); - for (const rawAssetId of returnedAssetIds) { - if (!isEvmAssetId(rawAssetId)) { - continue; - } - if (!hasPositiveBalance(returnedBalances[rawAssetId])) { - continue; - } - const normalizedAssetId = safeNormalize(rawAssetId); - if (!customSet.has(normalizedAssetId)) { - continue; - } - log('Graduating custom asset', { - accountId, - assetId: normalizedAssetId, - }); - this.#removeCustomAsset(accountId, normalizedAssetId); - } - - return next(ctx); - }); - } -} - -/** - * Check whether a CAIP-19 asset ID belongs to an EVM chain. - * - * @param assetId - The CAIP-19 asset ID to inspect. - * @returns `true` when the asset's chain namespace is `eip155`. - */ -function isEvmAssetId(assetId: Caip19AssetId): boolean { - // CAIP-19 format: :/: - // The chain namespace is always the segment before the first colon. - const namespace = assetId.split(':')[0]; - return namespace === KnownCaipNamespace.Eip155; -} - -/** - * Normalize a CAIP-19 asset ID, returning the original on failure. Some - * malformed IDs (e.g. an asset reference that fails address checksumming) - * make `normalizeAssetId` throw — in that case we fall back to the raw ID - * so the graduation pass can still proceed for other assets. - * - * @param assetId - The CAIP-19 asset ID to normalize. - * @returns The normalized ID, or the original on failure. - */ -function safeNormalize(assetId: Caip19AssetId): Caip19AssetId { - try { - return normalizeAssetId(assetId); - } catch { - return assetId; - } -} - -/** - * Whether a balance entry reports a strictly positive amount. AccountsAPI - * may return zero for tokens it indexes but the user no longer holds; we - * treat those as non-graduating so RPC keeps polling and surfaces any - * future incoming transfer immediately. - * - * `AssetBalance.amount` is already a human-readable decimal string from - * both AccountsApi (e.g. "0.283549083429656057") and the websocket data - * source (which divides by `decimals` before emitting), so a `Number()` - * sign check is safe: `NaN`, `undefined`, empty strings, and zero all - * fail the comparison. - * - * @param balance - The balance entry from the response. - * @returns `true` when the balance amount represents a value greater than 0. - */ -function hasPositiveBalance(balance: AssetBalance | undefined): boolean { - return Number(balance?.amount) > 0; -} diff --git a/packages/assets-controller/src/middlewares/ParallelMiddleware.ts b/packages/assets-controller/src/middlewares/ParallelMiddleware.ts index 8b6fceed491..d38fe358aad 100644 --- a/packages/assets-controller/src/middlewares/ParallelMiddleware.ts +++ b/packages/assets-controller/src/middlewares/ParallelMiddleware.ts @@ -53,6 +53,17 @@ export function mergeDataResponses(responses: DataResponse[]): DataResponse { ...response.errors, }; } + if ( + response.unprocessedCustomAssets && + response.unprocessedCustomAssets.length > 0 + ) { + merged.unprocessedCustomAssets = [ + ...new Set([ + ...(merged.unprocessedCustomAssets ?? []), + ...response.unprocessedCustomAssets, + ]), + ]; + } if (response.detectedAssets) { merged.detectedAssets = { ...(merged.detectedAssets ?? {}), diff --git a/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.test.ts b/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.test.ts index 4bcbbc472f4..9d6c276195e 100644 --- a/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.test.ts +++ b/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.test.ts @@ -14,6 +14,10 @@ const MOCK_ACCOUNT_ID = 'mock-account-id'; const MOCK_ASSET_MAINNET = 'eip155:1/slip44:60' as Caip19AssetId; const MOCK_ASSET_POLYGON = 'eip155:137/slip44:966' as Caip19AssetId; const MOCK_ASSET_BSC = 'eip155:56/slip44:714' as Caip19AssetId; +const MOCK_TOKEN_POLYGON = + 'eip155:137/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; +const MOCK_TOKEN_MAINNET = + 'eip155:1/erc20:0xdAC17F958D2ee523a2206206994597C13D831ec7' as Caip19AssetId; function createMockAccount(): InternalAccount { return { @@ -246,4 +250,115 @@ describe('RpcFallbackMiddleware', () => { const finalCtx = next.mock.calls[0][0]; expect(finalCtx.response.errors).toStrictEqual({}); }); + + it('recovers unprocessedCustomAssets with an RPC call scoped to just those assets (via customAssets), not a whole-chain fetch', async () => { + const rpcResponse: DataResponse = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [MOCK_TOKEN_POLYGON]: { amount: '3' } }, + }, + }; + const { source, middleware: rpcMw } = createMockRpcSource(rpcResponse); + const mw = new RpcFallbackMiddleware({ rpcDataSource: source }); + const ctx = createContext(createDataRequest(['eip155:1', 'eip155:137']), { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [MOCK_ASSET_MAINNET]: { amount: '1' } }, + }, + unprocessedCustomAssets: [MOCK_TOKEN_POLYGON], + }); + const next = jest.fn(async (innerCtx) => innerCtx); + + await mw.assetsMiddleware(ctx, next); + + expect(rpcMw).toHaveBeenCalledTimes(1); + const [rpcCtx] = rpcMw.mock.calls[0]; + expect(rpcCtx.request.chainIds).toStrictEqual(['eip155:137']); + expect(rpcCtx.request.customAssets).toStrictEqual([MOCK_TOKEN_POLYGON]); + + const finalCtx = next.mock.calls[0][0]; + expect(finalCtx.response.assetsBalance[MOCK_ACCOUNT_ID]).toStrictEqual({ + [MOCK_ASSET_MAINNET]: { amount: '1' }, + [MOCK_TOKEN_POLYGON]: { amount: '3' }, + }); + // Recovered — removed from the asset axis. + expect(finalCtx.response.unprocessedCustomAssets).toBeUndefined(); + }); + + it('skips asset-scoped recovery for assets whose chain was already retried on the chain axis', async () => { + // The chain-axis fetch (native + custom assets) already covers the token, so + // there must be no second, asset-scoped RPC call for the same chain. + const rpcResponse: DataResponse = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { [MOCK_TOKEN_POLYGON]: { amount: '3' } }, + }, + }; + const { source, middleware: rpcMw } = createMockRpcSource(rpcResponse); + const mw = new RpcFallbackMiddleware({ rpcDataSource: source }); + const ctx = createContext(createDataRequest(['eip155:137']), { + errors: { 'eip155:137': 'Unprocessed networks' }, + unprocessedCustomAssets: [MOCK_TOKEN_POLYGON], + }); + const next = jest.fn(async (innerCtx) => innerCtx); + + await mw.assetsMiddleware(ctx, next); + + // Only the chain-axis call — it uses the original request (no customAssets + // override to the unresolved subset). + expect(rpcMw).toHaveBeenCalledTimes(1); + const [rpcCtx] = rpcMw.mock.calls[0]; + expect(rpcCtx.request.chainIds).toStrictEqual(['eip155:137']); + expect(rpcCtx.request.customAssets).toBeUndefined(); + + const finalCtx = next.mock.calls[0][0]; + expect(finalCtx.response.errors).toStrictEqual({}); + expect(finalCtx.response.unprocessedCustomAssets).toBeUndefined(); + }); + + it('keeps unprocessedCustomAssets that RPC could not recover', async () => { + const { source } = createMockRpcSource({}); // RPC returns nothing + const mw = new RpcFallbackMiddleware({ rpcDataSource: source }); + const ctx = createContext(createDataRequest(['eip155:137']), { + unprocessedCustomAssets: [MOCK_TOKEN_POLYGON], + }); + const next = jest.fn(async (innerCtx) => innerCtx); + + await mw.assetsMiddleware(ctx, next); + + const finalCtx = next.mock.calls[0][0]; + expect(finalCtx.response.unprocessedCustomAssets).toStrictEqual([ + MOCK_TOKEN_POLYGON, + ]); + }); + + it('recovers both errored chains and unprocessed assets in separate RPC calls', async () => { + const rpcResponse: DataResponse = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_BSC]: { amount: '9' }, + [MOCK_TOKEN_MAINNET]: { amount: '4' }, + }, + }, + }; + const { source, middleware: rpcMw } = createMockRpcSource(rpcResponse); + const mw = new RpcFallbackMiddleware({ rpcDataSource: source }); + const ctx = createContext(createDataRequest(['eip155:1', 'eip155:56']), { + errors: { 'eip155:56': 'Fetch failed' }, + unprocessedCustomAssets: [MOCK_TOKEN_MAINNET], + }); + const next = jest.fn(async (innerCtx) => innerCtx); + + await mw.assetsMiddleware(ctx, next); + + // One call for the errored chain (whole chain), one scoped call for the pin. + expect(rpcMw).toHaveBeenCalledTimes(2); + const chainCall = rpcMw.mock.calls[0][0]; + expect(chainCall.request.chainIds).toStrictEqual(['eip155:56']); + expect(chainCall.request.customAssets).toBeUndefined(); + const assetCall = rpcMw.mock.calls[1][0]; + expect(assetCall.request.chainIds).toStrictEqual(['eip155:1']); + expect(assetCall.request.customAssets).toStrictEqual([MOCK_TOKEN_MAINNET]); + + const finalCtx = next.mock.calls[0][0]; + expect(finalCtx.response.errors).toStrictEqual({}); + expect(finalCtx.response.unprocessedCustomAssets).toBeUndefined(); + }); }); diff --git a/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.ts b/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.ts index fb164a89126..503dac31e81 100644 --- a/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.ts +++ b/packages/assets-controller/src/middlewares/RpcFallbackMiddleware.ts @@ -2,10 +2,14 @@ import { projectLogger, createModuleLogger } from '../logger.js'; import { forDataTypes } from '../types.js'; import type { AssetsDataSource, + Caip19AssetId, ChainId, + Context, + DataRequest, DataResponse, Middleware, } from '../types.js'; +import { normalizeAssetId } from '../utils/index.js'; import { mergeDataResponses } from './ParallelMiddleware.js'; const CONTROLLER_NAME = 'RpcFallbackMiddleware'; @@ -17,15 +21,20 @@ export type RpcFallbackMiddlewareOptions = { rpcDataSource: AssetsDataSource; }; +const noopNext = async (ctx: Context): Promise => ctx; + /** - * RpcFallbackMiddleware retries chains that failed upstream on the RPC data - * source. Any chain present in `response.errors` (network error, - * unprocessedNetworks, timeout, …) is handed off to RPC with the request - * filtered to just those chains. Successful RPC results are merged into the - * response and their entries are cleared from `response.errors`. + * RpcFallbackMiddleware recovers what upstream sources left outstanding on + * RPC, along two axes: + * + * - **Chain axis:** chains in `response.errors` are re-fetched in full and + * cleared from `response.errors` on recovery. + * - **Asset axis:** pins in `response.unprocessedCustomAssets` are re-fetched + * with an RPC request scoped to just those assets; recovered entries are + * pruned. Assets on chains already retried above are skipped. * - * Place this immediately after `createParallelBalanceMiddleware` in the fast - * pipeline. + * Place immediately after `createParallelBalanceMiddleware` in the fast and + * subscription enrichment pipelines. */ export class RpcFallbackMiddleware { readonly name = CONTROLLER_NAME; @@ -45,57 +54,154 @@ export class RpcFallbackMiddleware { const erroredChains = new Set( Object.keys(ctx.response.errors ?? {}) as ChainId[], ); - if (erroredChains.size === 0) { + const unprocessedCustomAssets = [ + ...new Set(ctx.response.unprocessedCustomAssets ?? []), + ]; + + if (erroredChains.size === 0 && unprocessedCustomAssets.length === 0) { return next(ctx); } - log('Retrying failed chains on RPC', { - chains: [...erroredChains], - }); - - const filteredRequest = { - ...ctx.request, - chainIds: ctx.request.chainIds.filter((id) => erroredChains.has(id)), - }; - - const noopNext = async (inner: typeof ctx): Promise => inner; - const rpcResult = await this.#rpcDataSource.assetsMiddleware( - { - ...ctx, - request: filteredRequest, - response: {}, - }, - noopNext, + let merged: DataResponse = ctx.response; + + // Chain axis: retry whole errored chains on RPC. + if (erroredChains.size > 0) { + merged = await this.#recoverErroredChains(ctx, merged, erroredChains); + } + + // Asset axis: recover unresolved pins; chains retried above already + // cover theirs. + const assetsToRecover = unprocessedCustomAssets.filter( + (assetId) => !erroredChains.has(chainIdOfAsset(assetId)), ); + if (assetsToRecover.length > 0) { + merged = await this.#recoverUnprocessedAssets( + ctx, + merged, + assetsToRecover, + ); + } + + // Prune asset-axis entries that now have a balance. + merged = clearRecoveredAssetIds(merged); + + return next({ ...ctx, response: merged }); + }); + } + + async #recoverErroredChains( + ctx: Context, + currentResponse: DataResponse, + erroredChains: Set, + ): Promise { + log('Retrying failed chains on RPC', { chains: [...erroredChains] }); + + const chainRequest: DataRequest = { + ...ctx.request, + chainIds: ctx.request.chainIds.filter((id) => erroredChains.has(id)), + }; + const rpcResult = await this.#rpcDataSource.assetsMiddleware( + { ...ctx, request: chainRequest, response: {} }, + noopNext, + ); + + const merged = mergeDataResponses([currentResponse, rpcResult.response]); - const merged: DataResponse = mergeDataResponses([ - ctx.response, - rpcResult.response, - ]); - - // Clear errors only for chains RPC actually recovered a balance for. - // We must inspect rpcResult.response — NOT merged — because merged - // also contains balances from the upstream sources (AccountsApi / - // Websocket / Staked). If those sources returned partial data for - // a chain that they also flagged as errored (e.g. via - // unprocessedNetworks), and RPC then failed for that same chain, - // looking at merged would incorrectly mark the error as recovered. - const rpcAssetsBalance = rpcResult.response.assetsBalance; - if (merged.errors && rpcAssetsBalance) { - const chainsRecoveredByRpc = new Set(); - for (const accountBalances of Object.values(rpcAssetsBalance)) { - for (const assetId of Object.keys(accountBalances)) { - chainsRecoveredByRpc.add(assetId.split('/')[0]); - } + // Clear errors only for chains RPC itself recovered. Inspect + // rpcResult.response — NOT merged — or partial upstream data for an + // errored chain would mark it recovered even when RPC failed. + const rpcAssetsBalance = rpcResult.response.assetsBalance; + if (merged.errors && rpcAssetsBalance) { + const chainsRecoveredByRpc = new Set(); + for (const accountBalances of Object.values(rpcAssetsBalance)) { + for (const assetId of Object.keys(accountBalances)) { + chainsRecoveredByRpc.add(assetId.split('/')[0]); } - for (const chainId of erroredChains) { - if (chainsRecoveredByRpc.has(chainId)) { - delete merged.errors[chainId]; - } + } + for (const chainId of erroredChains) { + if (chainsRecoveredByRpc.has(chainId)) { + delete merged.errors[chainId]; } } + } - return next({ ...ctx, response: merged }); + return merged; + } + + async #recoverUnprocessedAssets( + ctx: Context, + currentResponse: DataResponse, + assetsToRecover: Caip19AssetId[], + ): Promise { + const assetChains = [ + ...new Set(assetsToRecover.map((assetId) => chainIdOfAsset(assetId))), + ]; + + log('Recovering unprocessed pinned assets on RPC', { + assetIds: assetsToRecover, }); + + // Override `customAssets` so RPC fetches just the unresolved pins + // (native + pins in one multicall) instead of every pin on the chain. + const assetRequest: DataRequest = { + ...ctx.request, + chainIds: assetChains, + customAssets: assetsToRecover, + }; + const rpcResult = await this.#rpcDataSource.assetsMiddleware( + { ...ctx, request: assetRequest, response: {} }, + noopNext, + ); + + return mergeDataResponses([currentResponse, rpcResult.response]); + } +} + +/** + * Extract the CAIP-2 chain ID from a CAIP-19 asset ID. + * + * @param assetId - The CAIP-19 asset ID. + * @returns The CAIP-2 chain ID portion. + */ +function chainIdOfAsset(assetId: Caip19AssetId): ChainId { + return assetId.split('/')[0] as ChainId; +} + +/** + * Remove entries from `unprocessedCustomAssets` that now have a balance in + * the response. + * + * @param response - The merged data response. + * @returns The response with recovered assets pruned. + */ +function clearRecoveredAssetIds(response: DataResponse): DataResponse { + if ( + !response.unprocessedCustomAssets || + response.unprocessedCustomAssets.length === 0 + ) { + return response; + } + + const recovered = new Set(); + for (const accountBalances of Object.values(response.assetsBalance ?? {})) { + for (const assetId of Object.keys(accountBalances)) { + recovered.add(normalizeAssetId(assetId as Caip19AssetId)); + } + } + + const stillUnprocessed = response.unprocessedCustomAssets.filter( + (assetId) => !recovered.has(normalizeAssetId(assetId)), + ); + + if (stillUnprocessed.length === response.unprocessedCustomAssets.length) { + return response; + } + + const next = { ...response }; + if (stillUnprocessed.length === 0) { + delete next.unprocessedCustomAssets; + } else { + next.unprocessedCustomAssets = stillUnprocessed; } + return next; } diff --git a/packages/assets-controller/src/middlewares/index.ts b/packages/assets-controller/src/middlewares/index.ts index d6796382aa3..2cc36299fba 100644 --- a/packages/assets-controller/src/middlewares/index.ts +++ b/packages/assets-controller/src/middlewares/index.ts @@ -1,5 +1,3 @@ -export { CustomAssetGraduationMiddleware } from './CustomAssetGraduationMiddleware.js'; -export type { CustomAssetGraduationMiddlewareOptions } from './CustomAssetGraduationMiddleware.js'; export { DetectionMiddleware } from './DetectionMiddleware.js'; export { RpcFallbackMiddleware } from './RpcFallbackMiddleware.js'; export type { RpcFallbackMiddlewareOptions } from './RpcFallbackMiddleware.js'; diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index 2b2b940a713..6c4df24c4a3 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -339,12 +339,10 @@ export type DataRequest = { /** Specific CAIP-19 asset IDs */ customAssets?: Caip19AssetId[]; /** - * When true, the data source should poll only the user's `customAssets` - * for the requested chains and skip refreshing the regular tracked - * balances. Used by the AssetsController to issue a supplemental RPC - * subscription on chains that another data source is already covering. + * User-hidden CAIP-19 asset IDs, sent to the Accounts API v6 endpoint as + * `excludeAssetIds` so they are dropped from the response. */ - customAssetsOnly?: boolean; + excludeAssetIds?: Caip19AssetId[]; /** Force fresh fetch, bypass cache */ forceUpdate?: boolean; /** Hint for polling interval (ms) - used by data sources that implement polling */ @@ -363,8 +361,14 @@ export type DataResponse = { assetsPrice?: Record; /** Balance data per account */ assetsBalance?: Record>; - /** Errors encountered, keyed by chain ID */ + /** Errors encountered, keyed by chain ID (chain-axis fallback + telemetry) */ errors?: Record; + /** + * Pinned asset IDs the source could not resolve (asset-axis fallback). + * Unlike `errors` the chain succeeded — only these assets still need a + * downstream (RPC) fetch. + */ + unprocessedCustomAssets?: Caip19AssetId[]; /** Detected assets (assets that do not have metadata) */ detectedAssets?: Record; /** diff --git a/packages/assets-controller/src/utils/customAssetsRpcSupplement.test.ts b/packages/assets-controller/src/utils/customAssetsRpcSupplement.test.ts deleted file mode 100644 index 597a6e541b5..00000000000 --- a/packages/assets-controller/src/utils/customAssetsRpcSupplement.test.ts +++ /dev/null @@ -1,250 +0,0 @@ -import type { Caip19AssetId, ChainId } from '../types.js'; -import { pickRpcCustomAssetsSupplement } from './customAssetsRpcSupplement.js'; - -const MAINNET = 'eip155:1' as ChainId; -const POLYGON = 'eip155:137' as ChainId; -const OPTIMISM = 'eip155:10' as ChainId; -const SOLANA = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; - -const WETH_MAINNET = - 'eip155:1/erc20:0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2' as Caip19AssetId; -const USDC_MAINNET = - 'eip155:1/erc20:0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48' as Caip19AssetId; -const USDC_POLYGON = - 'eip155:137/erc20:0x3c499c542cef5e3811e1192ce70d8cc03d5c3359' as Caip19AssetId; -const TOKEN_OPTIMISM = - 'eip155:10/erc20:0x4200000000000000000000000000000000000042' as Caip19AssetId; -const SOLANA_TOKEN = - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB' as Caip19AssetId; - -const ACCOUNT_A = 'account-a'; -const ACCOUNT_B = 'account-b'; - -describe('pickRpcCustomAssetsSupplement', () => { - describe('the core invariant — customAssets must always be fetched by RPC', () => { - it('picks a chain claimed by AccountsApi/Websocket as long as the user has a customAsset there', () => { - // Mainnet has been claimed by another data source (AccountsApi or - // Websocket), so it is NOT in `rpcAssignedChains`. The user has a - // customAsset on mainnet — RPC must run a supplemental sub. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET, POLYGON]), - enabledChains: new Set([MAINNET, POLYGON]), - }); - - expect(result.chains).toStrictEqual([MAINNET]); - expect([...result.accountIds]).toStrictEqual([ACCOUNT_A]); - }); - - it('picks every chain that has a customAsset across all selected accounts', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A, ACCOUNT_B], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET, USDC_POLYGON], - [ACCOUNT_B]: [TOKEN_OPTIMISM], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET, POLYGON, OPTIMISM]), - enabledChains: new Set([MAINNET, POLYGON, OPTIMISM]), - }); - - expect(new Set(result.chains)).toStrictEqual( - new Set([MAINNET, POLYGON, OPTIMISM]), - ); - expect(result.accountIds).toStrictEqual(new Set([ACCOUNT_A, ACCOUNT_B])); - }); - - it('deduplicates a chain when multiple accounts have customAssets there', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A, ACCOUNT_B], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - [ACCOUNT_B]: [USDC_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([MAINNET]); - expect(result.accountIds).toStrictEqual(new Set([ACCOUNT_A, ACCOUNT_B])); - }); - }); - - describe('skip rules', () => { - it('skips a chain that the regular RPC subscription already covers', () => { - // The regular RPC sub already fetches customAssets for `MAINNET` - // (see BalanceFetcher#getAssetsToFetch), so a supplemental sub would - // double-poll. Skip. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set([MAINNET]), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([]); - }); - - it('skips a chain RPC cannot serve (no NetworkController config for it)', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set(), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([]); - }); - - it('skips a chain that is currently disabled', () => { - // The chain is RPC-supported and not RPC-assigned, but the user has - // disabled it — there is no UI surface that shows its balances, so - // polling would waste resources. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set(), - }); - - expect(result.chains).toStrictEqual([]); - }); - - it('skips an account with no customAssets', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A, ACCOUNT_B], - customAssetsByAccount: { - [ACCOUNT_A]: [], - [ACCOUNT_B]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([MAINNET]); - expect(result.accountIds).toStrictEqual(new Set([ACCOUNT_B])); - }); - - it('skips an account that is not in the selected accountIds list', () => { - // Even though state.customAssets has entries for ACCOUNT_B, the caller - // only selected ACCOUNT_A. ACCOUNT_B's customAssets are ignored. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [], - [ACCOUNT_B]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([]); - expect(result.accountIds).toStrictEqual(new Set()); - }); - - it('skips a malformed CAIP-19 asset ID without throwing', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: ['not-a-caip-19-id' as Caip19AssetId, WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([MAINNET]); - }); - }); - - describe('chain-namespace coverage', () => { - it('picks non-EVM chains too — the helper is namespace-agnostic', () => { - // The graduation middleware only graduates EVM customAssets, but the - // supplemental RPC fetch is conceptually independent: any chain RPC - // can serve and the user has imported a token on, gets supplemented. - // Solana is a hypothetical future case; today RPC reports it inactive - // so this test exists to document the invariant rather than gate - // production behavior. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [SOLANA_TOKEN], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([SOLANA]), - enabledChains: new Set([SOLANA]), - }); - - expect(result.chains).toStrictEqual([SOLANA]); - }); - }); - - describe('mixed scenarios that exercise multiple rules at once', () => { - it('picks one chain and skips another for the same account', () => { - // ACCOUNT_A holds: - // - WETH_MAINNET → RPC supplement picks it (claimed by API/WS). - // - USDC_POLYGON → regular RPC sub already covers it. Skip. - // - TOKEN_OPTIMISM → RPC doesn't support optimism here. Skip. - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET, USDC_POLYGON, TOKEN_OPTIMISM], - }, - rpcAssignedChains: new Set([POLYGON]), - rpcAvailableChains: new Set([MAINNET, POLYGON]), - enabledChains: new Set([MAINNET, POLYGON, OPTIMISM]), - }); - - expect(result.chains).toStrictEqual([MAINNET]); - expect(result.accountIds).toStrictEqual(new Set([ACCOUNT_A])); - }); - - it('returns empty when every customAsset is filtered out', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [ACCOUNT_A], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET, USDC_POLYGON], - }, - rpcAssignedChains: new Set([MAINNET, POLYGON]), - rpcAvailableChains: new Set([MAINNET, POLYGON]), - enabledChains: new Set([MAINNET, POLYGON]), - }); - - expect(result.chains).toStrictEqual([]); - // The account is still listed (the controller's empty-chains - // early-return prevents any subscription anyway). - expect(result.accountIds).toStrictEqual(new Set([ACCOUNT_A])); - }); - - it('returns empty when no accounts are passed in', () => { - const result = pickRpcCustomAssetsSupplement({ - accountIds: [], - customAssetsByAccount: { - [ACCOUNT_A]: [WETH_MAINNET], - }, - rpcAssignedChains: new Set(), - rpcAvailableChains: new Set([MAINNET]), - enabledChains: new Set([MAINNET]), - }); - - expect(result.chains).toStrictEqual([]); - expect(result.accountIds).toStrictEqual(new Set()); - }); - }); -}); diff --git a/packages/assets-controller/src/utils/customAssetsRpcSupplement.ts b/packages/assets-controller/src/utils/customAssetsRpcSupplement.ts deleted file mode 100644 index 3703a5f9fc9..00000000000 --- a/packages/assets-controller/src/utils/customAssetsRpcSupplement.ts +++ /dev/null @@ -1,112 +0,0 @@ -import { parseCaipAssetType } from '@metamask/utils'; - -import type { Caip19AssetId, ChainId } from '../types.js'; - -/** - * Inputs needed to decide whether `RpcDataSource` should run a supplemental - * `customAssetsOnly` subscription. All inputs are pre-computed by the caller - * (controller) so this helper is fully pure and deterministic. - */ -export type RpcCustomAssetsSupplementInput = { - /** Selected account IDs being considered for subscription. */ - accountIds: string[]; - /** `state.customAssets` slice — accountId → CAIP-19 asset IDs. */ - customAssetsByAccount: Record; - /** Chains the regular RPC subscription already covers. */ - rpcAssignedChains: Set; - /** Chains the RPC data source can serve (its current activeChains). */ - rpcAvailableChains: Set; - /** - * Chains that are currently enabled for at least one selected account - * (i.e. keys of the controller's chainToAccounts map). - */ - enabledChains: Set; -}; - -/** - * The decision output: which chains and accounts need a supplemental - * `customAssetsOnly` RPC subscription. - */ -export type RpcCustomAssetsSupplementResult = { - /** Chains that need a supplemental customAssetsOnly RPC subscription. */ - chains: ChainId[]; - /** Accounts that own at least one customAsset on a supplemental chain. */ - accountIds: Set; -}; - -/** - * Decide which chains require a supplemental `customAssetsOnly` RPC - * subscription so that user-imported customAssets are **always** polled by - * RPC — even when AccountsApi or the websocket data source has already - * claimed the chain in the regular handoff. - * - * RPC is the sole balance fetcher for customAssets: - * - AccountsApi indexes a curated token list and may return zero (or no - * entry at all) for a user-imported token, even on supported chains. - * - The websocket data source is push-only: it relays balance deltas from - * transactions, so a user who never transacts won't see balance updates. - * - * To guarantee freshness, we subscribe RPC in `customAssetsOnly` mode on - * any chain that: - * 1. has at least one selected account with a customAsset there; - * 2. is supported by the RPC data source; - * 3. is **not** already covered by the regular RPC subscription - * (`rpcAssignedChains`) — otherwise customAssets are picked up there; - * 4. is currently enabled (in `chainToAccounts`); polling a disabled - * chain wastes resources because the user can't see its balances. - * - * Malformed asset IDs are silently skipped — they can't be parsed into a - * chain anyway and shouldn't crash the subscription pipeline. - * - * @param input - Decision inputs. - * @returns Chains and account IDs that require supplemental subscription. - */ -export function pickRpcCustomAssetsSupplement( - input: RpcCustomAssetsSupplementInput, -): RpcCustomAssetsSupplementResult { - const { - accountIds, - customAssetsByAccount, - rpcAssignedChains, - rpcAvailableChains, - enabledChains, - } = input; - - const chains = new Set(); - const supplementalAccountIds = new Set(); - - for (const accountId of accountIds) { - const customForAccount = customAssetsByAccount[accountId]; - if (!customForAccount || customForAccount.length === 0) { - continue; - } - // Mirror the controller's prior behavior: an account with any customAsset - // (even if all are filtered out below) is added to the candidate set. - // The empty-chains early-return in the caller still prevents a useless - // subscription, so this preserves observable behavior. - supplementalAccountIds.add(accountId); - for (const assetId of customForAccount) { - let chainId: ChainId; - try { - chainId = parseCaipAssetType(assetId).chainId; - } catch { - continue; - } - if (rpcAssignedChains.has(chainId)) { - continue; - } - if (!rpcAvailableChains.has(chainId)) { - continue; - } - if (!enabledChains.has(chainId)) { - continue; - } - chains.add(chainId); - } - } - - return { - chains: [...chains], - accountIds: supplementalAccountIds, - }; -}