From e02faee23b7914754ccb8842c8642ad6d7ce63a8 Mon Sep 17 00:00:00 2001 From: Kriys94 Date: Wed, 1 Jul 2026 16:32:08 +0200 Subject: [PATCH 1/2] feat(assets-controller): use AccountsAPI v6 balances --- packages/assets-controller/CHANGELOG.md | 2 + packages/assets-controller/package.json | 1 + .../src/AssetsController.test.ts | 64 ++++ .../assets-controller/src/AssetsController.ts | 6 +- .../AccountsApiDataSource.test.ts | 335 +++++++++++++++++- .../src/data-sources/AccountsApiDataSource.ts | 295 +++++++++++++-- .../assets-controller/tsconfig.build.json | 3 +- packages/assets-controller/tsconfig.json | 3 +- yarn.lock | 1 + 9 files changed, 675 insertions(+), 35 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 3e0a758c56..14894443e5 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Add Robinhood Chain (`4663`/`0x1237`) in `MulticallClient` ([#9443](https://github.com/MetaMask/core/pull/9443)) +- `AccountsApiDataSource` now selects the Accounts API balances endpoint version from the `RemoteFeatureFlagController` (`assetsAccountsApiV6` flag, read per fetch off the shared messenger, default v5) so the v6 endpoint is gated consistently across clients (extension, mobile) without each client wiring a getter. Adds a required `messenger` option and `RemoteFeatureFlagController:getState` to `AccountsApiDataSourceAllowedActions`. Only `category: 'token'` rows from the v6 response are consumed (DeFi positions are ignored) to preserve parity with v5 ([#9344](https://github.com/MetaMask/core/pull/9344)) +- Add `AccountsApiDataSourceConfig.useBalanceV6` optional getter (`() => boolean`) that overrides the remote-flag-driven endpoint selection (e.g. for tests); when omitted the `assetsAccountsApiV6` remote flag is used ([#9344](https://github.com/MetaMask/core/pull/9344)) ### Changed diff --git a/packages/assets-controller/package.json b/packages/assets-controller/package.json index 2fed58b489..558c42ceb6 100644 --- a/packages/assets-controller/package.json +++ b/packages/assets-controller/package.json @@ -74,6 +74,7 @@ "@metamask/phishing-controller": "^17.2.0", "@metamask/polling-controller": "^16.0.8", "@metamask/preferences-controller": "^23.1.0", + "@metamask/remote-feature-flag-controller": "^4.2.2", "@metamask/snaps-controllers": "^19.0.0", "@metamask/snaps-utils": "^12.1.2", "@metamask/transaction-controller": "^68.3.0", diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 8fbd85f741..4ccd5c021c 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -18,6 +18,7 @@ import type { AssetsControllerMessenger, AssetsControllerState, } from './AssetsController'; +import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource'; import type { PriceDataSourceConfig } from './data-sources/PriceDataSource'; import { PriceDataSource } from './data-sources/PriceDataSource'; import { TokenDataSource } from './data-sources/TokenDataSource'; @@ -121,10 +122,16 @@ type WithControllerOptions = { * Required for tests that rely on asset tracking running (e.g. trace on unlock). */ clientControllerState?: { isUiOpen: boolean }; + /** + * When set, registers RemoteFeatureFlagController:getState so the controller can + * read feature flags (e.g. `assetsAccountsApiV6` gating the balances endpoint). + */ + remoteFeatureFlags?: Record; /** Extra options passed to AssetsController constructor (e.g. trace). */ controllerOptions?: Partial<{ trace: TraceCallback; priceDataSourceConfig: PriceDataSourceConfig; + accountsApiDataSourceConfig: AccountsApiDataSourceConfig; isEnabled: () => boolean; captureException: (error: Error) => void; tempMigrateAssetsInfoMetadataAssets3346: () => Assets3346MigrationState; @@ -156,6 +163,7 @@ async function withController( state = {}, isBasicFunctionality = (): boolean => true, clientControllerState, + remoteFeatureFlags, queryApiClient = createMockQueryApiClient(), controllerOptions = {}, }, @@ -230,6 +238,17 @@ async function withController( ); } + if (remoteFeatureFlags !== undefined) { + ( + messenger as { + registerActionHandler: (a: string, h: () => unknown) => void; + } + ).registerActionHandler('RemoteFeatureFlagController:getState', () => ({ + remoteFeatureFlags, + cacheTimestamp: 0, + })); + } + const controller = new AssetsController({ messenger: messenger as unknown as AssetsControllerMessenger, state, @@ -878,6 +897,51 @@ describe('AssetsController', () => { }); }); + // Endpoint selection from the flag is unit-tested in AccountsApiDataSource; + // this asserts the controller wires its messenger through so the + // `assetsAccountsApiV6` flag drives endpoint selection end-to-end. + it('routes to the Accounts API v6 endpoint when the assetsAccountsApiV6 remote flag is enabled', async () => { + const fetchV6MultiAccountBalances = jest.fn().mockResolvedValue({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + }); + const fetchV5MultiAccountBalances = jest.fn().mockResolvedValue({ + balances: [], + unprocessedNetworks: [], + }); + + const queryApiClient = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: [1], + partialSupport: [], + }), + fetchV6MultiAccountBalances, + fetchV5MultiAccountBalances, + }, + } as unknown as ApiPlatformClient; + + await withController( + { + queryApiClient, + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }, + async ({ controller }) => { + await flushPromises(); + + await controller.getAssets([createMockInternalAccount()], { + chainIds: ['eip155:1'], + forceUpdate: true, + }); + + expect(fetchV6MultiAccountBalances).toHaveBeenCalled(); + expect(fetchV5MultiAccountBalances).not.toHaveBeenCalled(); + }, + ); + }); + describe('pipeline splitting', () => { it('returns from getAssets before background pipelines complete', async () => { // Spy on handleAssetsUpdate to count how many times state is written. diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 3a6612af09..081856f6b0 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -48,6 +48,7 @@ import type { } from '@metamask/permission-controller'; import { PhishingControllerBulkScanTokensAction } from '@metamask/phishing-controller'; import type { PreferencesControllerStateChangeEvent } from '@metamask/preferences-controller'; +import type { RemoteFeatureFlagControllerGetStateAction } from '@metamask/remote-feature-flag-controller'; import type { SnapControllerGetRunnableSnapsAction, SnapControllerHandleRequestAction, @@ -321,7 +322,9 @@ type AllowedActions = // BackendWebsocketDataSource | BackendWebSocketServiceActions // PhishingController - | PhishingControllerBulkScanTokensAction; + | PhishingControllerBulkScanTokensAction + // AccountsApiDataSource (Accounts API v6 balances feature flag) + | RemoteFeatureFlagControllerGetStateAction; type AccountTreeControllerStateChangedEvent = ControllerStateChangedEvent< 'AccountTreeController', @@ -914,6 +917,7 @@ export class AssetsController extends BaseController< this.#getAssetType(assetId), }); this.#accountsApiDataSource = new AccountsApiDataSource({ + messenger: this.messenger, queryApiClient, onActiveChainsUpdated: this.#onActiveChainsUpdated, ...accountsApiDataSourceConfig, diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index 645b0a44de..6d0b4d58e6 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -1,11 +1,16 @@ /* eslint-disable jest/unbound-method */ -import type { V5BalanceItem } from '@metamask/core-backend'; +import type { + V5BalanceItem, + V6BalanceItem, + V6AccountBalancesEntry, +} from '@metamask/core-backend'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; import type { MockAnyNamespace } from '@metamask/messenger'; import type { ChainId, + Caip19AssetId, DataRequest, Context, AssetsControllerStateInternal, @@ -32,6 +37,7 @@ type MockApiClient = { accounts: { fetchV2SupportedNetworks: jest.Mock; fetchV5MultiAccountBalances: jest.Mock; + fetchV6MultiAccountBalances: jest.Mock; }; }; @@ -59,6 +65,7 @@ function createMockApiClient( supportedChains: number[] = [1, 137], balances: V5BalanceItem[] = [], unprocessedNetworks: string[] = [], + v6Accounts: V6AccountBalancesEntry[] = [], ): MockApiClient { return { accounts: { @@ -70,10 +77,23 @@ function createMockApiClient( balances, unprocessedNetworks, }), + fetchV6MultiAccountBalances: jest.fn().mockResolvedValue({ + accounts: v6Accounts, + unprocessedNetworks, + unprocessedIncludeAssetIds: [], + }), }, }; } +function createMockV6BalanceItem( + assetId: string, + balance: string, + category: 'token' | 'defi' = 'token', +): V6BalanceItem { + return { category, assetId, balance } as V6BalanceItem; +} + function createMockBalanceItem( accountId: string, assetId: string, @@ -122,6 +142,9 @@ async function setupController( balances?: V5BalanceItem[]; unprocessedNetworks?: string[]; fetchTimeoutMs?: number; + v6Accounts?: V6AccountBalancesEntry[]; + useBalanceV6?: () => boolean; + remoteFeatureFlags?: Record; } = {}, ): Promise { const { @@ -129,6 +152,9 @@ async function setupController( balances = [], unprocessedNetworks = [], fetchTimeoutMs, + v6Accounts = [], + useBalanceV6, + remoteFeatureFlags, } = options; const rootMessenger = new Messenger({ @@ -145,9 +171,23 @@ async function setupController( parent: rootMessenger, }); + if (remoteFeatureFlags !== undefined) { + ( + rootMessenger as unknown as { + registerActionHandler: (a: string, h: () => unknown) => void; + } + ).registerActionHandler('RemoteFeatureFlagController:getState', () => ({ + remoteFeatureFlags, + cacheTimestamp: 0, + })); + } + rootMessenger.delegate({ messenger: controllerMessenger, - actions: [], + actions: + remoteFeatureFlags === undefined + ? [] + : ['RemoteFeatureFlagController:getState'], events: [], }); @@ -158,14 +198,18 @@ async function setupController( supportedChains, balances, unprocessedNetworks, + v6Accounts, ); const controller = new AccountsApiDataSource({ + messenger: + controllerMessenger as unknown as AccountsApiDataSourceOptions['messenger'], queryApiClient: apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'], onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), ...(fetchTimeoutMs === undefined ? {} : { fetchTimeoutMs }), + ...(useBalanceV6 === undefined ? {} : { useBalanceV6 }), }); // Wait for async initialization @@ -323,7 +367,9 @@ describe('AccountsApiDataSource', () => { }); it('fetch calls API with correct account IDs', async () => { - const { controller, apiClient } = await setupController(); + const { controller, apiClient } = await setupController({ + useBalanceV6: () => false, + }); await controller.fetch(createDataRequest()); @@ -337,7 +383,9 @@ describe('AccountsApiDataSource', () => { }); it('fetch bypasses TanStack cache when forceUpdate is true', async () => { - const { controller, apiClient } = await setupController(); + const { controller, apiClient } = await setupController({ + useBalanceV6: () => false, + }); await controller.fetch(createDataRequest({ forceUpdate: true })); @@ -359,7 +407,10 @@ describe('AccountsApiDataSource', () => { ), ]; - const { controller } = await setupController({ balances }); + const { controller } = await setupController({ + balances, + useBalanceV6: () => false, + }); const response = await controller.fetch(createDataRequest()); @@ -389,7 +440,9 @@ describe('AccountsApiDataSource', () => { }); it('fetch handles API errors', async () => { - const { controller, apiClient } = await setupController(); + const { controller, apiClient } = await setupController({ + useBalanceV6: () => false, + }); apiClient.accounts.fetchV5MultiAccountBalances.mockRejectedValueOnce( new Error('API Error'), @@ -402,9 +455,271 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); + describe('useBalanceV6 feature flag', () => { + it('uses the v5 endpoint by default', async () => { + const { controller, apiClient } = await setupController(); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).not.toHaveBeenCalled(); + + controller.destroy(); + }); + + it('uses the v6 endpoint when the assetsAccountsApiV6 remote flag is enabled', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: true }, + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).not.toHaveBeenCalled(); + + controller.destroy(); + }); + + it('uses the v5 endpoint when the assetsAccountsApiV6 remote flag is disabled', async () => { + const { controller, apiClient } = await setupController({ + remoteFeatureFlags: { assetsAccountsApiV6: false }, + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).not.toHaveBeenCalled(); + + controller.destroy(); + }); + + it('uses the v5 endpoint when the flag returns false', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => false, + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).not.toHaveBeenCalled(); + + controller.destroy(); + }); + + it('uses the v6 endpoint when the flag returns true', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => true, + }); + + await controller.fetch(createDataRequest()); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith([`eip155:1:${MOCK_ADDRESS}`], undefined, undefined); + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).not.toHaveBeenCalled(); + + controller.destroy(); + }); + + it('reads the flag per fetch so it can revert to v5 at runtime', async () => { + let v6Enabled = true; + const { controller, apiClient } = await setupController({ + useBalanceV6: () => v6Enabled, + }); + + await controller.fetch(createDataRequest()); + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + + v6Enabled = false; + await controller.fetch(createDataRequest()); + expect( + apiClient.accounts.fetchV5MultiAccountBalances, + ).toHaveBeenCalledTimes(1); + + controller.destroy(); + }); + + it('processes v6 token balances grouped by account', async () => { + const { controller } = await setupController({ + useBalanceV6: () => true, + v6Accounts: [ + { + accountId: `eip155:1:${MOCK_ADDRESS}`, + balances: [ + createMockV6BalanceItem( + 'eip155:1/slip44:60', + '1000000000000000000', + ), + ], + }, + ], + }); + + const response = await controller.fetch(createDataRequest()); + + expect( + response.assetsBalance?.['mock-account-id']?.['eip155:1/slip44:60'] + ?.amount, + ).toBe('1000000000000000000'); + + controller.destroy(); + }); + + it('ignores v6 defi positions', async () => { + const { controller } = await setupController({ + useBalanceV6: () => true, + v6Accounts: [ + { + accountId: `eip155:1:${MOCK_ADDRESS}`, + balances: [ + createMockV6BalanceItem( + 'eip155:1/slip44:60', + '1000000000000000000', + ), + createMockV6BalanceItem( + 'eip155:1/erc20:0xdefi', + '500', + 'defi', + ), + ], + }, + ], + }); + + const response = await controller.fetch(createDataRequest()); + + const accountBalances = response.assetsBalance?.['mock-account-id'] ?? {}; + expect(accountBalances).toHaveProperty('eip155:1/slip44:60'); + expect(accountBalances).not.toHaveProperty('eip155:1/erc20:0xdefi'); + + controller.destroy(); + }); + + it('marks v6 unprocessed networks as errors', async () => { + const { controller } = await setupController({ + useBalanceV6: () => true, + unprocessedNetworks: ['eip155:1'], + }); + + const response = await controller.fetch(createDataRequest()); + + expect(response.errors?.[CHAIN_MAINNET]).toBe( + 'Unprocessed by Accounts API', + ); + + controller.destroy(); + }); + + it('handles v6 API errors', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => true, + }); + + apiClient.accounts.fetchV6MultiAccountBalances.mockRejectedValueOnce( + new Error('API Error'), + ); + + const response = await controller.fetch(createDataRequest()); + + expect(response.errors?.[CHAIN_MAINNET]).toContain('Fetch failed'); + + controller.destroy(); + }); + + it('passes custom ERC-20 tokens to v6 as includeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => true, + }); + + const customToken = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ customAssets: [customToken] }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [customToken] }, + undefined, + ); + + controller.destroy(); + }); + + it('filters out native and non-erc20 custom assets from includeAssetIds', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => true, + }); + + const erc20Token = + 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; + const nativeAsset = 'eip155:1/slip44:60' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ customAssets: [erc20Token, nativeAsset] }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + { includeAssetIds: [erc20Token] }, + undefined, + ); + + controller.destroy(); + }); + + it('excludes custom tokens on chains not part of the fetch', async () => { + const { controller, apiClient } = await setupController({ + useBalanceV6: () => true, + }); + + const otherChainToken = + 'eip155:137/erc20:0x0000000000000000000000000000000000001010' as Caip19AssetId; + + await controller.fetch( + createDataRequest({ + chainIds: [CHAIN_MAINNET], + customAssets: [otherChainToken], + }), + ); + + expect( + apiClient.accounts.fetchV6MultiAccountBalances, + ).toHaveBeenCalledWith([`eip155:1:${MOCK_ADDRESS}`], undefined, undefined); + + controller.destroy(); + }); + }); + it('fetch marks every requested chain as errored when the call exceeds the configured timeout', async () => { const { controller, apiClient } = await setupController({ fetchTimeoutMs: 10, + useBalanceV6: () => false, }); apiClient.accounts.fetchV5MultiAccountBalances.mockImplementationOnce( @@ -462,7 +777,10 @@ describe('AccountsApiDataSource', () => { ), ]; - const { controller } = await setupController({ balances }); + const { controller } = await setupController({ + balances, + useBalanceV6: () => false, + }); const next = jest.fn().mockResolvedValue(undefined); const context = createMiddlewareContext(); @@ -577,10 +895,13 @@ describe('AccountsApiDataSource', () => { ); const controllerOptions: AccountsApiDataSourceOptions = { + messenger: + controllerMessenger as unknown as AccountsApiDataSourceOptions['messenger'], queryApiClient: apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'], onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), + useBalanceV6: (): boolean => false, }; if (tokenDetectionEnabled !== undefined) { diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index 74f5d94be1..99259ee08e 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -1,11 +1,17 @@ -import type { V5BalanceItem } from '@metamask/core-backend'; +import type { + V5BalanceItem, + V6AccountBalancesEntry, +} from '@metamask/core-backend'; import { ApiPlatformClient } from '@metamask/core-backend'; +import type { RemoteFeatureFlagControllerGetStateAction } from '@metamask/remote-feature-flag-controller'; import { isCaipChainId, KnownCaipNamespace, + parseCaipAssetType, toCaipChainId, } from '@metamask/utils'; +import type { AssetsControllerMessenger } from '../AssetsController'; import { projectLogger, createModuleLogger } from '../logger'; import type { ChainId, @@ -37,9 +43,11 @@ const log = createModuleLogger(projectLogger, CONTROLLER_NAME); // MESSENGER TYPES // ============================================================================ -// Allowed actions that AccountsApiDataSource can call (none - uses callbacks). -// Note: Uses ApiPlatformClient directly, so no BackendApiClient actions needed -export type AccountsApiDataSourceAllowedActions = never; +// Allowed actions that AccountsApiDataSource can call. Balances are fetched via +// ApiPlatformClient directly (no BackendApiClient actions needed); the messenger +// is only used to read the Accounts API v6 balances feature flag. +export type AccountsApiDataSourceAllowedActions = + RemoteFeatureFlagControllerGetStateAction; // ============================================================================ // STATE @@ -71,9 +79,25 @@ export type AccountsApiDataSourceConfig = { * middleware hands them off to the next data source (e.g. RPC fallback). */ fetchTimeoutMs?: number; + /** + * Optional override for the Accounts API balances endpoint version. + * + * When omitted, the data source reads the `assetsAccountsApiV6` remote feature + * flag from the RemoteFeatureFlagController on every fetch (defaulting to v5 + * when the flag is unset or unavailable). Provide a getter to bypass the flag + * (e.g. in tests). When it returns true the v6 endpoint is used; otherwise the + * legacy v5 endpoint is used. Kept as a getter so the value can be toggled at + * runtime without re-instantiating the data source. + */ + useBalanceV6?: () => boolean; }; export type AccountsApiDataSourceOptions = AccountsApiDataSourceConfig & { + /** + * The AssetsController messenger (shared by all data sources). Used to read + * the `assetsAccountsApiV6` remote feature flag. + */ + messenger: AssetsControllerMessenger; /** ApiPlatformClient for API calls with caching */ queryApiClient: ApiPlatformClient; /** Called when active chains are updated. Pass dataSourceName so the controller knows the source. */ @@ -192,6 +216,16 @@ export class AccountsApiDataSource extends AbstractDataSource< /** Getter avoids stale value when user toggles token detection at runtime. */ readonly #tokenDetectionEnabled: () => boolean; + /** + * Feature flag getter selecting the balances endpoint version. Kept as a + * getter so it can be flipped at runtime (e.g. to revert v6 -> v5). Defaults + * to reading the `assetsAccountsApiV6` remote feature flag off the messenger. + */ + readonly #useBalanceV6: () => boolean; + + /** Shared AssetsController messenger, used to read remote feature flags. */ + readonly #messenger: AssetsControllerMessenger; + /** ApiPlatformClient for cached API calls */ readonly #apiClient: ApiPlatformClient; @@ -212,11 +246,34 @@ export class AccountsApiDataSource extends AbstractDataSource< this.#fetchTimeoutMs = options.fetchTimeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS; this.#tokenDetectionEnabled = options.tokenDetectionEnabled ?? ((): boolean => true); + this.#messenger = options.messenger; + this.#useBalanceV6 = + options.useBalanceV6 ?? ((): boolean => this.#isBalanceV6Enabled()); this.#apiClient = options.queryApiClient; this.#initializeActiveChains().catch(console.error); } + /** + * Whether the Accounts API v6 balances endpoint is enabled, read from the + * RemoteFeatureFlagController (`assetsAccountsApiV6`). Read on demand (per + * fetch) rather than cached so the value can be toggled at runtime, and behaves + * the same across clients (extension, mobile). Defaults to `false` when the + * flag is unset or the controller is unavailable. + * + * @returns `true` when the v6 balances endpoint should be used. + */ + #isBalanceV6Enabled(): boolean { + try { + const { remoteFeatureFlags } = this.#messenger.call( + 'RemoteFeatureFlagController:getState', + ); + return remoteFeatureFlags?.assetsAccountsApiV6 === true; + } catch { + return false; + } + } + // ============================================================================ // INITIALIZATION // ============================================================================ @@ -334,20 +391,20 @@ export class AccountsApiDataSource extends AbstractDataSource< ? { staleTime: 0, gcTime: 0 } : undefined; - const apiResponse = await fetchWithTimeout( - () => - this.#apiClient.accounts.fetchV5MultiAccountBalances( + // 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.#useBalanceV6() + ? await this.#fetchV6Balances( accountIds, - undefined, fetchOptions, - ), - this.#fetchTimeoutMs, - ); + request, + this.#getIncludeAssetIds(request, chainsToFetch), + ) + : await this.#fetchV5Balances(accountIds, fetchOptions, request); // Handle unprocessed networks - these will be passed to next middleware - if (apiResponse.unprocessedNetworks.length > 0) { - const unprocessedChainIds = - apiResponse.unprocessedNetworks.map(caipChainIdToChainId); + if (unprocessedNetworks.length > 0) { + const unprocessedChainIds = unprocessedNetworks.map(caipChainIdToChainId); // Add unprocessed chains to errors so middleware passes them to next data source response.errors = response.errors ?? {}; @@ -356,11 +413,6 @@ export class AccountsApiDataSource extends AbstractDataSource< } } - const { assetsBalance } = this.#processV5Balances( - apiResponse.balances, - request, - ); - response.assetsBalance = assetsBalance; response.updateMode = 'merge'; } catch (error) { @@ -390,6 +442,140 @@ export class AccountsApiDataSource extends AbstractDataSource< return response; } + /** + * Fetch balances from the legacy v5 endpoint and process them. + * + * @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. + */ + async #fetchV5Balances( + accountIds: string[], + fetchOptions: { staleTime: number; gcTime: number } | undefined, + request: DataRequest, + ): Promise<{ + unprocessedNetworks: string[]; + assetsBalance: Record>; + }> { + const apiResponse = await fetchWithTimeout( + () => + this.#apiClient.accounts.fetchV5MultiAccountBalances( + accountIds, + undefined, + fetchOptions, + ), + this.#fetchTimeoutMs, + ); + + const { assetsBalance } = this.#processV5Balances( + apiResponse.balances, + request, + ); + + return { + unprocessedNetworks: apiResponse.unprocessedNetworks, + assetsBalance, + }; + } + + /** + * Fetch balances from the v6 endpoint and process them. + * + * @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. + * @param includeAssetIds - Custom ERC-20 CAIP-19 asset IDs to confirm + * detection for (passed to the v6 endpoint as `includeAssetIds`). + * @returns Unprocessed networks and processed asset balances by account. + */ + async #fetchV6Balances( + accountIds: string[], + fetchOptions: { staleTime: number; gcTime: number } | undefined, + request: DataRequest, + includeAssetIds: string[], + ): Promise<{ + unprocessedNetworks: string[]; + assetsBalance: Record>; + }> { + const apiResponse = await fetchWithTimeout( + () => + this.#apiClient.accounts.fetchV6MultiAccountBalances( + accountIds, + includeAssetIds.length > 0 ? { includeAssetIds } : undefined, + fetchOptions, + ), + this.#fetchTimeoutMs, + ); + + const { assetsBalance } = this.#processV6Balances( + apiResponse.accounts, + request, + ); + + return { + unprocessedNetworks: apiResponse.unprocessedNetworks, + assetsBalance, + }; + } + + /** + * Derive the v6 `includeAssetIds` list from the request's custom assets. + * + * The v6 balances endpoint only accepts ERC-20 CAIP-19 asset IDs here, so + * native (`slip44`) and non-EVM custom assets are filtered out. Assets on + * chains not part of this fetch are also dropped. + * + * @param request - The original data request containing custom assets. + * @param chainsToFetch - Chains being fetched in this request. + * @returns Deduplicated ERC-20 CAIP-19 asset IDs to include. + */ + #getIncludeAssetIds( + request: DataRequest, + chainsToFetch: ChainId[], + ): string[] { + if (!request.customAssets || request.customAssets.length === 0) { + return []; + } + + const chainSet = new Set(chainsToFetch); + const includeAssetIds = new Set(); + + for (const assetId of request.customAssets) { + try { + const parsed = parseCaipAssetType(assetId); + // v6 includeAssetIds only supports ERC-20 tokens. + if (parsed.assetNamespace !== 'erc20') { + continue; + } + const chainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; + if (chainSet.has(chainId)) { + includeAssetIds.add(assetId); + } + } catch { + // Skip unparseable asset IDs + } + } + + return [...includeAssetIds]; + } + + /** + * Build a lookup of lowercased account address to the request's account ID. + * + * @param request - The original data request containing accounts to map. + * @returns Map of lowercase address to account ID. + */ + #buildAddressToAccountIdMap(request: DataRequest): Map { + const addressToAccountId = new Map(); + for (const { account } of request.accountsWithSupportedChains) { + if (account.address) { + addressToAccountId.set(account.address.toLowerCase(), account.id); + } + } + return addressToAccountId; + } + /** * Process V5 API balances response. * V5 returns a flat array of balance items, each with accountId and assetId. @@ -410,12 +596,7 @@ export class AccountsApiDataSource extends AbstractDataSource< > = {}; // Build a map of lowercase addresses to account IDs for efficient lookup - const addressToAccountId = new Map(); - for (const { account } of request.accountsWithSupportedChains) { - if (account.address) { - addressToAccountId.set(account.address.toLowerCase(), account.id); - } - } + const addressToAccountId = this.#buildAddressToAccountIdMap(request); // V5 response: array of { accountId, assetId, balance, ... } for (const item of balances) { @@ -449,6 +630,70 @@ export class AccountsApiDataSource extends AbstractDataSource< return { assetsBalance }; } + /** + * Process V6 API balances response. + * V6 groups balances per account (`accounts: [{ accountId, balances }]`). + * Only `category: 'token'` rows are consumed here to preserve parity with + * the v5 token-balance behavior; DeFi positions are ignored. + * + * @param accounts - Per-account balance entries from the V6 API response. + * @param request - The original data request containing accounts to map. + * @returns Object containing processed asset balances by account. + */ + #processV6Balances( + accounts: V6AccountBalancesEntry[], + request: DataRequest, + ): { + assetsBalance: Record>; + } { + const assetsBalance: Record< + string, + Record + > = {}; + + // Build a map of lowercase addresses to account IDs for efficient lookup + const addressToAccountId = this.#buildAddressToAccountIdMap(request); + + for (const entry of accounts) { + // Extract address from CAIP-10 account ID (e.g., "eip155:1:0x1234..." -> "0x1234...") + const addressParts = entry.accountId.split(':'); + if (addressParts.length < 3) { + continue; + } + const address = addressParts[2].toLowerCase(); + + // Find the matching account ID from request + const accountId = addressToAccountId.get(address); + if (!accountId) { + // This is normal - API returns balances for all chains, but request may only have one account + continue; + } + + for (const item of entry.balances) { + // Only consume token balances; DeFi positions are handled elsewhere. + if (item.category !== 'token') { + continue; + } + + if (!assetsBalance[accountId]) { + assetsBalance[accountId] = {}; + } + + // Normalize asset ID (checksum EVM addresses for ERC20 tokens) + const normalizedAssetId = normalizeAssetId( + item.assetId as Caip19AssetId, + ); + + // Store balance as returned by API + assetsBalance[accountId][normalizedAssetId] = { + amount: item.balance, + }; + } + } + + return { assetsBalance }; + } + // ============================================================================ // MIDDLEWARE // ============================================================================ diff --git a/packages/assets-controller/tsconfig.build.json b/packages/assets-controller/tsconfig.build.json index f7468098e5..4454b4442b 100644 --- a/packages/assets-controller/tsconfig.build.json +++ b/packages/assets-controller/tsconfig.build.json @@ -16,7 +16,8 @@ { "path": "../network-enablement-controller/tsconfig.build.json" }, { "path": "../permission-controller/tsconfig.build.json" }, { "path": "../phishing-controller/tsconfig.build.json" }, - { "path": "../preferences-controller/tsconfig.build.json" } + { "path": "../preferences-controller/tsconfig.build.json" }, + { "path": "../remote-feature-flag-controller/tsconfig.build.json" } ], "include": ["../../types", "./src"], "exclude": ["**/*.test.ts", "**/__fixtures__/"] diff --git a/packages/assets-controller/tsconfig.json b/packages/assets-controller/tsconfig.json index aa015e2729..731de50f36 100644 --- a/packages/assets-controller/tsconfig.json +++ b/packages/assets-controller/tsconfig.json @@ -13,7 +13,8 @@ { "path": "../messenger" }, { "path": "../network-enablement-controller" }, { "path": "../phishing-controller" }, - { "path": "../preferences-controller" } + { "path": "../preferences-controller" }, + { "path": "../remote-feature-flag-controller" } ], "include": ["../../types", "./src"] } diff --git a/yarn.lock b/yarn.lock index e44825fb42..035eb8ec65 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5686,6 +5686,7 @@ __metadata: "@metamask/phishing-controller": "npm:^17.2.0" "@metamask/polling-controller": "npm:^16.0.8" "@metamask/preferences-controller": "npm:^23.1.0" + "@metamask/remote-feature-flag-controller": "npm:^4.2.2" "@metamask/snaps-controllers": "npm:^19.0.0" "@metamask/snaps-utils": "npm:^12.1.2" "@metamask/transaction-controller": "npm:^68.3.0" From ed27e0c6a2d33fd0e2ac07c6f5c7579060ce5b4f Mon Sep 17 00:00:00 2001 From: Christian Tran Date: Thu, 9 Jul 2026 11:34:58 +0200 Subject: [PATCH 2/2] Potential fix for pull request finding 'CodeQL / Prototype-polluting assignment' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- README.md | 1 + packages/assets-controller/CHANGELOG.md | 3 +- .../src/AssetsController.test.ts | 2 +- .../AccountsApiDataSource.test.ts | 111 +++++------------- .../src/data-sources/AccountsApiDataSource.ts | 99 ++++------------ 5 files changed, 51 insertions(+), 165 deletions(-) diff --git a/README.md b/README.md index a528c55398..be341858b8 100644 --- a/README.md +++ b/README.md @@ -281,6 +281,7 @@ linkStyle default opacity:0.5 assets_controller --> phishing_controller; assets_controller --> polling_controller; assets_controller --> preferences_controller; + assets_controller --> remote_feature_flag_controller; assets_controller --> transaction_controller; assets_controllers --> account_tree_controller; assets_controllers --> accounts_controller; diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 14894443e5..6f04635e67 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -10,8 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Add Robinhood Chain (`4663`/`0x1237`) in `MulticallClient` ([#9443](https://github.com/MetaMask/core/pull/9443)) -- `AccountsApiDataSource` now selects the Accounts API balances endpoint version from the `RemoteFeatureFlagController` (`assetsAccountsApiV6` flag, read per fetch off the shared messenger, default v5) so the v6 endpoint is gated consistently across clients (extension, mobile) without each client wiring a getter. Adds a required `messenger` option and `RemoteFeatureFlagController:getState` to `AccountsApiDataSourceAllowedActions`. Only `category: 'token'` rows from the v6 response are consumed (DeFi positions are ignored) to preserve parity with v5 ([#9344](https://github.com/MetaMask/core/pull/9344)) -- Add `AccountsApiDataSourceConfig.useBalanceV6` optional getter (`() => boolean`) that overrides the remote-flag-driven endpoint selection (e.g. for tests); when omitted the `assetsAccountsApiV6` remote flag is used ([#9344](https://github.com/MetaMask/core/pull/9344)) +- `AccountsApiDataSource` now selects the Accounts API balances endpoint version from the `RemoteFeatureFlagController` (`assetsAccountsApiV6` flag, read per fetch off the shared messenger, default v5) so the v6 endpoint is gated consistently across clients (extension, mobile) without each client wiring a getter. The flag is read as a JSON variation shaped `{ value: boolean }` (same shape as `backendWebSocketConnection`). Adds a required `messenger` option and `RemoteFeatureFlagController:getState` to `AccountsApiDataSourceAllowedActions`. Only `category: 'token'` rows from the v6 response are consumed (DeFi positions are ignored) to preserve parity with v5 ([#9344](https://github.com/MetaMask/core/pull/9344)) ### Changed diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 4ccd5c021c..0adc63a2b7 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -926,7 +926,7 @@ describe('AssetsController', () => { await withController( { queryApiClient, - remoteFeatureFlags: { assetsAccountsApiV6: true }, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }, async ({ controller }) => { await flushPromises(); diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts index 6d0b4d58e6..f01829bf15 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.test.ts @@ -143,7 +143,6 @@ async function setupController( unprocessedNetworks?: string[]; fetchTimeoutMs?: number; v6Accounts?: V6AccountBalancesEntry[]; - useBalanceV6?: () => boolean; remoteFeatureFlags?: Record; } = {}, ): Promise { @@ -153,7 +152,6 @@ async function setupController( unprocessedNetworks = [], fetchTimeoutMs, v6Accounts = [], - useBalanceV6, remoteFeatureFlags, } = options; @@ -209,7 +207,6 @@ async function setupController( onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), ...(fetchTimeoutMs === undefined ? {} : { fetchTimeoutMs }), - ...(useBalanceV6 === undefined ? {} : { useBalanceV6 }), }); // Wait for async initialization @@ -367,9 +364,7 @@ describe('AccountsApiDataSource', () => { }); it('fetch calls API with correct account IDs', async () => { - const { controller, apiClient } = await setupController({ - useBalanceV6: () => false, - }); + const { controller, apiClient } = await setupController(); await controller.fetch(createDataRequest()); @@ -383,9 +378,7 @@ describe('AccountsApiDataSource', () => { }); it('fetch bypasses TanStack cache when forceUpdate is true', async () => { - const { controller, apiClient } = await setupController({ - useBalanceV6: () => false, - }); + const { controller, apiClient } = await setupController(); await controller.fetch(createDataRequest({ forceUpdate: true })); @@ -409,7 +402,6 @@ describe('AccountsApiDataSource', () => { const { controller } = await setupController({ balances, - useBalanceV6: () => false, }); const response = await controller.fetch(createDataRequest()); @@ -440,9 +432,7 @@ describe('AccountsApiDataSource', () => { }); it('fetch handles API errors', async () => { - const { controller, apiClient } = await setupController({ - useBalanceV6: () => false, - }); + const { controller, apiClient } = await setupController(); apiClient.accounts.fetchV5MultiAccountBalances.mockRejectedValueOnce( new Error('API Error'), @@ -455,7 +445,7 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - describe('useBalanceV6 feature flag', () => { + describe('assetsAccountsApiV6 feature flag', () => { it('uses the v5 endpoint by default', async () => { const { controller, apiClient } = await setupController(); @@ -473,7 +463,7 @@ describe('AccountsApiDataSource', () => { it('uses the v6 endpoint when the assetsAccountsApiV6 remote flag is enabled', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: true }, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }); await controller.fetch(createDataRequest()); @@ -490,7 +480,7 @@ describe('AccountsApiDataSource', () => { it('uses the v5 endpoint when the assetsAccountsApiV6 remote flag is disabled', async () => { const { controller, apiClient } = await setupController({ - remoteFeatureFlags: { assetsAccountsApiV6: false }, + remoteFeatureFlags: { assetsAccountsApiV6: { value: false } }, }); await controller.fetch(createDataRequest()); @@ -505,9 +495,9 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('uses the v5 endpoint when the flag returns false', async () => { + it('uses the v5 endpoint when the flag is a plain boolean instead of the JSON value shape', async () => { const { controller, apiClient } = await setupController({ - useBalanceV6: () => false, + remoteFeatureFlags: { assetsAccountsApiV6: true }, }); await controller.fetch(createDataRequest()); @@ -522,16 +512,20 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('uses the v6 endpoint when the flag returns true', async () => { + it('calls the v6 endpoint without extra params when enabled', async () => { const { controller, apiClient } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }); await controller.fetch(createDataRequest()); expect( apiClient.accounts.fetchV6MultiAccountBalances, - ).toHaveBeenCalledWith([`eip155:1:${MOCK_ADDRESS}`], undefined, undefined); + ).toHaveBeenCalledWith( + [`eip155:1:${MOCK_ADDRESS}`], + undefined, + undefined, + ); expect( apiClient.accounts.fetchV5MultiAccountBalances, ).not.toHaveBeenCalled(); @@ -540,9 +534,11 @@ describe('AccountsApiDataSource', () => { }); it('reads the flag per fetch so it can revert to v5 at runtime', async () => { - let v6Enabled = true; + const remoteFeatureFlags = { + assetsAccountsApiV6: { value: true }, + }; const { controller, apiClient } = await setupController({ - useBalanceV6: () => v6Enabled, + remoteFeatureFlags, }); await controller.fetch(createDataRequest()); @@ -550,7 +546,7 @@ describe('AccountsApiDataSource', () => { apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledTimes(1); - v6Enabled = false; + remoteFeatureFlags.assetsAccountsApiV6 = { value: false }; await controller.fetch(createDataRequest()); expect( apiClient.accounts.fetchV5MultiAccountBalances, @@ -561,7 +557,7 @@ describe('AccountsApiDataSource', () => { it('processes v6 token balances grouped by account', async () => { const { controller } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, v6Accounts: [ { accountId: `eip155:1:${MOCK_ADDRESS}`, @@ -587,7 +583,7 @@ describe('AccountsApiDataSource', () => { it('ignores v6 defi positions', async () => { const { controller } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, v6Accounts: [ { accountId: `eip155:1:${MOCK_ADDRESS}`, @@ -596,11 +592,7 @@ describe('AccountsApiDataSource', () => { 'eip155:1/slip44:60', '1000000000000000000', ), - createMockV6BalanceItem( - 'eip155:1/erc20:0xdefi', - '500', - 'defi', - ), + createMockV6BalanceItem('eip155:1/erc20:0xdefi', '500', 'defi'), ], }, ], @@ -617,7 +609,7 @@ describe('AccountsApiDataSource', () => { it('marks v6 unprocessed networks as errors', async () => { const { controller } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, unprocessedNetworks: ['eip155:1'], }); @@ -632,7 +624,7 @@ describe('AccountsApiDataSource', () => { it('handles v6 API errors', async () => { const { controller, apiClient } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }); apiClient.accounts.fetchV6MultiAccountBalances.mockRejectedValueOnce( @@ -646,9 +638,9 @@ describe('AccountsApiDataSource', () => { controller.destroy(); }); - it('passes custom ERC-20 tokens to v6 as includeAssetIds', async () => { + it('does not pass includeAssetIds to v6 even when custom assets are present', async () => { const { controller, apiClient } = await setupController({ - useBalanceV6: () => true, + remoteFeatureFlags: { assetsAccountsApiV6: { value: true } }, }); const customToken = @@ -662,64 +654,17 @@ describe('AccountsApiDataSource', () => { apiClient.accounts.fetchV6MultiAccountBalances, ).toHaveBeenCalledWith( [`eip155:1:${MOCK_ADDRESS}`], - { includeAssetIds: [customToken] }, undefined, - ); - - controller.destroy(); - }); - - it('filters out native and non-erc20 custom assets from includeAssetIds', async () => { - const { controller, apiClient } = await setupController({ - useBalanceV6: () => true, - }); - - const erc20Token = - 'eip155:1/erc20:0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48' as Caip19AssetId; - const nativeAsset = 'eip155:1/slip44:60' as Caip19AssetId; - - await controller.fetch( - createDataRequest({ customAssets: [erc20Token, nativeAsset] }), - ); - - expect( - apiClient.accounts.fetchV6MultiAccountBalances, - ).toHaveBeenCalledWith( - [`eip155:1:${MOCK_ADDRESS}`], - { includeAssetIds: [erc20Token] }, undefined, ); controller.destroy(); }); - - it('excludes custom tokens on chains not part of the fetch', async () => { - const { controller, apiClient } = await setupController({ - useBalanceV6: () => true, - }); - - const otherChainToken = - 'eip155:137/erc20:0x0000000000000000000000000000000000001010' as Caip19AssetId; - - await controller.fetch( - createDataRequest({ - chainIds: [CHAIN_MAINNET], - customAssets: [otherChainToken], - }), - ); - - expect( - apiClient.accounts.fetchV6MultiAccountBalances, - ).toHaveBeenCalledWith([`eip155:1:${MOCK_ADDRESS}`], undefined, undefined); - - controller.destroy(); - }); }); it('fetch marks every requested chain as errored when the call exceeds the configured timeout', async () => { const { controller, apiClient } = await setupController({ fetchTimeoutMs: 10, - useBalanceV6: () => false, }); apiClient.accounts.fetchV5MultiAccountBalances.mockImplementationOnce( @@ -779,7 +724,6 @@ describe('AccountsApiDataSource', () => { const { controller } = await setupController({ balances, - useBalanceV6: () => false, }); const next = jest.fn().mockResolvedValue(undefined); @@ -901,7 +845,6 @@ describe('AccountsApiDataSource', () => { apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'], onActiveChainsUpdated: (dataSourceName, chains, previousChains): void => activeChainsUpdateHandler(dataSourceName, chains, previousChains), - useBalanceV6: (): boolean => false, }; if (tokenDetectionEnabled !== undefined) { diff --git a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts index 99259ee08e..6ac6167b0e 100644 --- a/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts +++ b/packages/assets-controller/src/data-sources/AccountsApiDataSource.ts @@ -7,7 +7,6 @@ import type { RemoteFeatureFlagControllerGetStateAction } from '@metamask/remote import { isCaipChainId, KnownCaipNamespace, - parseCaipAssetType, toCaipChainId, } from '@metamask/utils'; @@ -79,17 +78,6 @@ export type AccountsApiDataSourceConfig = { * middleware hands them off to the next data source (e.g. RPC fallback). */ fetchTimeoutMs?: number; - /** - * Optional override for the Accounts API balances endpoint version. - * - * When omitted, the data source reads the `assetsAccountsApiV6` remote feature - * flag from the RemoteFeatureFlagController on every fetch (defaulting to v5 - * when the flag is unset or unavailable). Provide a getter to bypass the flag - * (e.g. in tests). When it returns true the v6 endpoint is used; otherwise the - * legacy v5 endpoint is used. Kept as a getter so the value can be toggled at - * runtime without re-instantiating the data source. - */ - useBalanceV6?: () => boolean; }; export type AccountsApiDataSourceOptions = AccountsApiDataSourceConfig & { @@ -216,13 +204,6 @@ export class AccountsApiDataSource extends AbstractDataSource< /** Getter avoids stale value when user toggles token detection at runtime. */ readonly #tokenDetectionEnabled: () => boolean; - /** - * Feature flag getter selecting the balances endpoint version. Kept as a - * getter so it can be flipped at runtime (e.g. to revert v6 -> v5). Defaults - * to reading the `assetsAccountsApiV6` remote feature flag off the messenger. - */ - readonly #useBalanceV6: () => boolean; - /** Shared AssetsController messenger, used to read remote feature flags. */ readonly #messenger: AssetsControllerMessenger; @@ -247,8 +228,6 @@ export class AccountsApiDataSource extends AbstractDataSource< this.#tokenDetectionEnabled = options.tokenDetectionEnabled ?? ((): boolean => true); this.#messenger = options.messenger; - this.#useBalanceV6 = - options.useBalanceV6 ?? ((): boolean => this.#isBalanceV6Enabled()); this.#apiClient = options.queryApiClient; this.#initializeActiveChains().catch(console.error); @@ -261,6 +240,10 @@ export class AccountsApiDataSource extends AbstractDataSource< * the same across clients (extension, mobile). Defaults to `false` when the * flag is unset or the controller is unavailable. * + * The flag is a LaunchDarkly JSON variation shaped `{ value: boolean }` (same + * shape as `backendWebSocketConnection`), so the nested `value` is read rather + * than treating the flag as a plain boolean. + * * @returns `true` when the v6 balances endpoint should be used. */ #isBalanceV6Enabled(): boolean { @@ -268,7 +251,12 @@ export class AccountsApiDataSource extends AbstractDataSource< const { remoteFeatureFlags } = this.#messenger.call( 'RemoteFeatureFlagController:getState', ); - return remoteFeatureFlags?.assetsAccountsApiV6 === true; + const flag = remoteFeatureFlags?.assetsAccountsApiV6; + return ( + typeof flag === 'object' && + flag !== null && + Boolean((flag as { value?: unknown }).value) + ); } catch { return false; } @@ -393,18 +381,14 @@ export class AccountsApiDataSource extends AbstractDataSource< // 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.#useBalanceV6() - ? await this.#fetchV6Balances( - accountIds, - fetchOptions, - request, - this.#getIncludeAssetIds(request, chainsToFetch), - ) + const { unprocessedNetworks, assetsBalance } = this.#isBalanceV6Enabled() + ? await this.#fetchV6Balances(accountIds, fetchOptions, request) : await this.#fetchV5Balances(accountIds, fetchOptions, request); // Handle unprocessed networks - these will be passed to next middleware if (unprocessedNetworks.length > 0) { - const unprocessedChainIds = unprocessedNetworks.map(caipChainIdToChainId); + const unprocessedChainIds = + unprocessedNetworks.map(caipChainIdToChainId); // Add unprocessed chains to errors so middleware passes them to next data source response.errors = response.errors ?? {}; @@ -485,15 +469,12 @@ 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. - * @param includeAssetIds - Custom ERC-20 CAIP-19 asset IDs to confirm - * detection for (passed to the v6 endpoint as `includeAssetIds`). * @returns Unprocessed networks and processed asset balances by account. */ async #fetchV6Balances( accountIds: string[], fetchOptions: { staleTime: number; gcTime: number } | undefined, request: DataRequest, - includeAssetIds: string[], ): Promise<{ unprocessedNetworks: string[]; assetsBalance: Record>; @@ -502,7 +483,7 @@ export class AccountsApiDataSource extends AbstractDataSource< () => this.#apiClient.accounts.fetchV6MultiAccountBalances( accountIds, - includeAssetIds.length > 0 ? { includeAssetIds } : undefined, + undefined, fetchOptions, ), this.#fetchTimeoutMs, @@ -519,47 +500,6 @@ export class AccountsApiDataSource extends AbstractDataSource< }; } - /** - * Derive the v6 `includeAssetIds` list from the request's custom assets. - * - * The v6 balances endpoint only accepts ERC-20 CAIP-19 asset IDs here, so - * native (`slip44`) and non-EVM custom assets are filtered out. Assets on - * chains not part of this fetch are also dropped. - * - * @param request - The original data request containing custom assets. - * @param chainsToFetch - Chains being fetched in this request. - * @returns Deduplicated ERC-20 CAIP-19 asset IDs to include. - */ - #getIncludeAssetIds( - request: DataRequest, - chainsToFetch: ChainId[], - ): string[] { - if (!request.customAssets || request.customAssets.length === 0) { - return []; - } - - const chainSet = new Set(chainsToFetch); - const includeAssetIds = new Set(); - - for (const assetId of request.customAssets) { - try { - const parsed = parseCaipAssetType(assetId); - // v6 includeAssetIds only supports ERC-20 tokens. - if (parsed.assetNamespace !== 'erc20') { - continue; - } - const chainId = `${parsed.chain.namespace}:${parsed.chain.reference}`; - if (chainSet.has(chainId)) { - includeAssetIds.add(assetId); - } - } catch { - // Skip unparseable asset IDs - } - } - - return [...includeAssetIds]; - } - /** * Build a lookup of lowercased account address to the request's account ID. * @@ -646,10 +586,10 @@ export class AccountsApiDataSource extends AbstractDataSource< ): { assetsBalance: Record>; } { - const assetsBalance: Record< + const assetsBalance = Object.create(null) as Record< string, Record - > = {}; + >; // Build a map of lowercase addresses to account IDs for efficient lookup const addressToAccountId = this.#buildAddressToAccountIdMap(request); @@ -676,7 +616,10 @@ export class AccountsApiDataSource extends AbstractDataSource< } if (!assetsBalance[accountId]) { - assetsBalance[accountId] = {}; + assetsBalance[accountId] = Object.create(null) as Record< + Caip19AssetId, + AssetBalance + >; } // Normalize asset ID (checksum EVM addresses for ERC20 tokens)