From 14e36eb402a8a044dba8997a1f989113f5ae6720 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Tue, 7 Jul 2026 17:06:19 +0200 Subject: [PATCH 1/9] chore: asset enrichment through snap data source --- packages/assets-controller/CHANGELOG.md | 5 + packages/assets-controller/package.json | 2 +- .../src/AssetsController.test.ts | 34 ++ .../assets-controller/src/AssetsController.ts | 9 +- .../src/data-sources/SnapDataSource.test.ts | 322 +++++++++++++++++- .../src/data-sources/SnapDataSource.ts | 253 ++++++++++++-- packages/assets-controller/src/types.ts | 19 ++ 7 files changed, 618 insertions(+), 26 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index cf2afd1720..d437a721f9 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,12 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add optional snap account-asset enrichment to fungible balances on `SnapDataSource.fetch` so Stellar trustline fields (`accountAssetInfo`) are included when clients call `getAssets`; balance push events (`AccountsController:accountBalancesUpdated`) pass amounts through without enrichment — clients should call `getAssets` after trustline-related snap events (e.g. `AccountsController:accountAssetListUpdated`) + ### Changed - Bump `@metamask/messenger` from `^1.2.0` to `^2.0.0` ([#9392](https://github.com/MetaMask/core/pull/9392)) ### Fixed +- Refresh Stellar trustline enrichment for custom assets that drop off `listAccountAssets` after deactivate (e.g. limit 0 tombstones) by including `customAssets` in `SnapDataSource.fetch` balance and enrichment requests - Fetch balances when switching account groups, enabling RPC-only networks, or after a new account is added to the account tree ([#9388](https://github.com/MetaMask/core/pull/9388)) ## [10.0.1] diff --git a/packages/assets-controller/package.json b/packages/assets-controller/package.json index 9fa1ae7755..834aa386ca 100644 --- a/packages/assets-controller/package.json +++ b/packages/assets-controller/package.json @@ -1,6 +1,6 @@ { "name": "@metamask/assets-controller", - "version": "10.0.1", + "version": "10.0.1-dev.8", "description": "Tracks assets balances/prices and handles token detection across all digital assets", "keywords": [ "Ethereum", diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 3222a8d285..9442987fd9 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -1687,6 +1687,40 @@ describe('AssetsController', () => { }); }); + it('preserves accountAssetInfo enrichment when merge update changes amount', async () => { + const initialState: Partial = { + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { + amount: '1', + accountAssetInfo: { limit: '1000', authorized: true }, + }, + }, + }, + }; + + await withController({ state: initialState }, async ({ controller }) => { + await controller.handleAssetsUpdate( + { + updateMode: 'merge', + assetsBalance: { + [MOCK_ACCOUNT_ID]: { + [MOCK_ASSET_ID]: { amount: '2' }, + }, + }, + }, + 'TestSource', + ); + + expect( + controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], + ).toStrictEqual({ + amount: '2', + accountAssetInfo: { limit: '1000', authorized: true }, + }); + }); + }); + it('preserves existing balances when merge update adds new chain data', async () => { const polygonNative = 'eip155:137/slip44:966' as Caip19AssetId; const initialState: Partial = { diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 9e99e94d7c..8e558d7c2c 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -574,7 +574,14 @@ function mergeAccountBalances( replaceCoveredChains: boolean, ): Record { if (!replaceCoveredChains) { - return { ...previousBalances, ...accountBalances }; + const next: Record = { ...previousBalances }; + for (const [assetId, balance] of Object.entries(accountBalances)) { + next[assetId] = { + ...(previousBalances[assetId] ?? { amount: '0' }), + ...balance, + }; + } + return next; } const coveredChains = new Set( diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index 8518b3d694..791ee402ff 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -28,10 +28,12 @@ import { const SOLANA_MAINNET = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; const BITCOIN_MAINNET = 'bip122:000000000019d6689c085ae165831e93' as ChainId; const TRON_MAINNET = 'tron:728126428' as ChainId; +const STELLAR_PUBNET = 'stellar:pubnet' as ChainId; // Test snap IDs const SOLANA_SNAP_ID = 'npm:@metamask/solana-wallet-snap'; const BITCOIN_SNAP_ID = 'npm:@metamask/bitcoin-wallet-snap'; +const STELLAR_SNAP_ID = 'npm:@metamask/stellar-wallet-snap'; type AllActions = SnapDataSourceAllowedActions; type AllEvents = SnapDataSourceAllowedEvents; @@ -43,6 +45,10 @@ const MOCK_SOL_ASSET = const MOCK_BTC_ASSET = 'bip122:000000000019d6689c085ae165831e93/slip44:0' as Caip19AssetId; const MOCK_TRON_ASSET = 'tron:728126428/slip44:195' as Caip19AssetId; +const MOCK_STELLAR_USDC_ASSET = + 'stellar:pubnet/asset:USDC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; +const MOCK_STELLAR_AUDD_ASSET = + 'stellar:pubnet/asset:AUDD-GDC7X2MXTYSAKUUGAIQ7J7RPEIM7GXSAIWFYWWH4GLNFECQVJJLB2EEU' as Caip19AssetId; const CHAIN_MAINNET = 'eip155:1' as ChainId; type SetupResult = { @@ -156,6 +162,7 @@ function createMockPermissions( function createMockHandleRequest( accountAssets: string[] = [], balances: Record = {}, + accountAssetInfo: Record = {}, ): jest.Mock { return jest.fn().mockImplementation((params) => { const { request } = params; @@ -165,6 +172,9 @@ function createMockHandleRequest( if (request?.method === 'keyring_getAccountBalances') { return Promise.resolve(balances); } + if (request?.method === 'getAccountAssetInfo') { + return Promise.resolve(accountAssetInfo); + } return Promise.resolve(null); }); } @@ -174,10 +184,16 @@ function setupController( installedSnaps?: Record; accountAssets?: string[]; balances?: Record; + accountAssetInfo?: Record; configuredNetworks?: ChainId[]; } = {}, ): SetupResult { - const { installedSnaps = {}, accountAssets = [], balances = {} } = options; + const { + installedSnaps = {}, + accountAssets = [], + balances = {}, + accountAssetInfo = {}, + } = options; const rootMessenger = new Messenger({ namespace: MOCK_ANY_NAMESPACE, @@ -226,7 +242,11 @@ function setupController( mockGetRunnableSnaps, ); - const mockHandleRequest = createMockHandleRequest(accountAssets, balances); + const mockHandleRequest = createMockHandleRequest( + accountAssets, + balances, + accountAssetInfo, + ); rootMessenger.registerActionHandler( 'SnapController:handleRequest', mockHandleRequest, @@ -521,6 +541,260 @@ describe('SnapDataSource', () => { cleanup(); }); + it('fetch enriches Stellar assets with account asset info', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssets: [MOCK_STELLAR_USDC_ASSET], + balances: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, + }, + accountAssetInfo: { + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + sponsored: false, + }, + }, + }); + await new Promise(process.nextTick); + + const response = await controller.fetch( + createDataRequest({ + chainIds: [STELLAR_PUBNET], + accounts: [ + createMockAccount({ + scopes: [STELLAR_PUBNET], + metadata: { + name: 'Stellar Account', + importTime: Date.now(), + keyring: { type: 'Snap Keyring' }, + snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, + }, + }), + ], + }), + ); + + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ + snapId: STELLAR_SNAP_ID, + origin: 'metamask', + handler: 'onClientRequest', + request: expect.objectContaining({ + method: 'getAccountAssetInfo', + params: { + accountId: 'mock-account-id', + scope: STELLAR_PUBNET, + assets: [MOCK_STELLAR_USDC_ASSET], + }, + }), + }), + ); + expect( + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], + ).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '1000', + authorized: true, + sponsored: false, + }, + }); + + cleanup(); + }); + + it('fetch enriches custom Stellar assets missing from listAccountAssets', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssets: [MOCK_STELLAR_USDC_ASSET], + balances: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, + [MOCK_STELLAR_AUDD_ASSET]: { amount: '0', unit: 'AUDD' }, + }, + accountAssetInfo: { + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + }, + [MOCK_STELLAR_AUDD_ASSET]: { + limit: '0', + authorized: true, + }, + }, + }); + await new Promise(process.nextTick); + + const response = await controller.fetch( + createDataRequest({ + chainIds: [STELLAR_PUBNET], + customAssets: [MOCK_STELLAR_AUDD_ASSET], + accounts: [ + createMockAccount({ + scopes: [STELLAR_PUBNET], + metadata: { + name: 'Stellar Account', + importTime: Date.now(), + keyring: { type: 'Snap Keyring' }, + snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, + }, + }), + ], + }), + ); + + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + method: 'keyring_getAccountBalances', + params: { + id: 'mock-account-id', + assets: expect.arrayContaining([ + MOCK_STELLAR_USDC_ASSET, + MOCK_STELLAR_AUDD_ASSET, + ]), + }, + }), + }), + ); + expect( + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_AUDD_ASSET], + ).toStrictEqual({ + amount: '0', + accountAssetInfo: { + limit: '0', + authorized: true, + }, + }); + + cleanup(); + }); + + it('fetch returns balance without accountAssetInfo when snap hangs past timeout', async () => { + jest.useFakeTimers(); + + const stellarAccount = createMockAccount({ + scopes: [STELLAR_PUBNET], + metadata: { + name: 'Stellar Account', + importTime: Date.now(), + keyring: { type: 'Snap Keyring' }, + snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, + }, + }); + + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssets: [MOCK_STELLAR_USDC_ASSET], + balances: { [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' } }, + }); + + // Drain all microtasks from subscribe's initial fetch before replacing the mock + await jest.advanceTimersByTimeAsync(20_000); + + // Now configure getAccountAssetInfo to never resolve + mockHandleRequest.mockImplementation((params) => { + const { request } = params as { request: { method: string } }; + if (request?.method === 'keyring_listAccountAssets') { + return Promise.resolve([MOCK_STELLAR_USDC_ASSET]); + } + if (request?.method === 'keyring_getAccountBalances') { + return Promise.resolve({ + [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, + }); + } + if (request?.method === 'getAccountAssetInfo') { + return new Promise(() => undefined); // never resolves + } + return Promise.resolve(null); + }); + + const fetchPromise = controller.fetch( + createDataRequest({ chainIds: [STELLAR_PUBNET], accounts: [stellarAccount] }), + ); + + // Advance past the enrichment timeout + await jest.advanceTimersByTimeAsync(20_000); + const response = await fetchPromise; + + expect( + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], + ).toStrictEqual({ amount: '25' }); + expect( + (response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET] as Record)?.accountAssetInfo, + ).toBeUndefined(); + + cleanup(); + jest.useRealTimers(); + }); + + it('fetch stops queuing batches after first batch times out', async () => { + jest.useFakeTimers(); + + const MOCK_STELLAR_ASSET_2 = + 'stellar:pubnet/asset:USDT-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; + const allAssets = [MOCK_STELLAR_USDC_ASSET, MOCK_STELLAR_ASSET_2]; + const balances = Object.fromEntries( + allAssets.map((a) => [a, { amount: '10', unit: 'X' }]), + ) as Record; + + const stellarAccount = createMockAccount({ + scopes: [STELLAR_PUBNET], + metadata: { + name: 'Stellar Account', + importTime: Date.now(), + keyring: { type: 'Snap Keyring' }, + snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, + }, + }); + + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssets: allAssets, + balances, + }); + + // Drain timers from subscribe's initial fetch + await jest.advanceTimersByTimeAsync(20_000); + + let getAccountAssetInfoCallCount = 0; + mockHandleRequest.mockImplementation((params) => { + const { request } = params as { request: { method: string } }; + if (request?.method === 'keyring_listAccountAssets') { + return Promise.resolve(allAssets); + } + if (request?.method === 'keyring_getAccountBalances') { + return Promise.resolve(balances); + } + if (request?.method === 'getAccountAssetInfo') { + getAccountAssetInfoCallCount += 1; + return new Promise(() => undefined); // always hangs + } + return Promise.resolve(null); + }); + + const fetchPromise = controller.fetch( + createDataRequest({ chainIds: [STELLAR_PUBNET], accounts: [stellarAccount] }), + ); + + await jest.advanceTimersByTimeAsync(20_000); + await fetchPromise; + + // Only 1 getAccountAssetInfo call: batch loop must break on timeout, not continue + expect(getAccountAssetInfoCallCount).toBe(1); + + cleanup(); + jest.useRealTimers(); + }); + it('fetch handles empty account assets gracefully', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ installedSnaps: { @@ -600,6 +874,50 @@ describe('SnapDataSource', () => { cleanup(); }); + it('does not enrich Stellar assets from snap balances updated event', async () => { + const { triggerBalancesUpdated, assetsUpdateHandler, mockHandleRequest, cleanup } = + setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssetInfo: { + [MOCK_STELLAR_USDC_ASSET]: { limit: '500' }, + }, + }); + await new Promise(process.nextTick); + + triggerBalancesUpdated({ + balances: { + 'account-1': { + [MOCK_STELLAR_USDC_ASSET]: { amount: '5', unit: 'USDC' }, + }, + }, + }); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(assetsUpdateHandler).toHaveBeenCalledWith( + expect.objectContaining({ + assetsBalance: { + 'account-1': { + [MOCK_STELLAR_USDC_ASSET]: { + amount: '5', + }, + }, + }, + }), + ); + expect(mockHandleRequest).not.toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + method: 'getAccountAssetInfo', + }), + }), + ); + + cleanup(); + }); + it('filters assets for chains without discovered snaps from balance update event', async () => { const { triggerBalancesUpdated, assetsUpdateHandler, cleanup } = setupController({ diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index 17da7b0b5e..e0db593540 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -25,6 +25,7 @@ import type { Caip19AssetId, DataRequest, DataResponse, + GetAccountAssetInfoResponse, Middleware, } from '../types'; import { AbstractDataSource } from './AbstractDataSource'; @@ -74,6 +75,28 @@ export const KEYRING_PERMISSION = 'endowment:keyring'; /** The permission name for snap assets endowment (contains chainIds) */ export const ASSETS_PERMISSION = 'endowment:assets'; +/** Snap clientRequest method for per-(account, asset) enrichment data. */ +const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; + +/** + * Max assets per snap getAccountAssetInfo request. Large batches can terminate + * the Stellar wallet snap on mobile when many trustlines are fetched at once. + */ +const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 3; + +/** Per-batch snap client request timeout (ms). Hung requests must not block apply. */ +const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; + +/** + * Chains whose wallet snap implements getAccountAssetInfo. + */ +const ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN: Partial< + Record +> = { + 'stellar:pubnet': true, + 'stellar:testnet': true, +}; + // ============================================================================ // PERMISSION UTILITIES // ============================================================================ @@ -115,6 +138,10 @@ export function extractChainFromAssetId(assetId: string): ChainId { return parsed.chainId; } +function isAccountAssetInfoEnrichmentAvailable(chainId: ChainId): boolean { + return ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN[chainId] === true; +} + // ============================================================================ // STATE // ============================================================================ @@ -266,6 +293,11 @@ export class SnapDataSource extends AbstractDataSource< * Handle snap balance updates from the keyring. * Transforms the payload and publishes to AssetsController. * + * Push updates carry amounts only. Per-asset snap enrichment (e.g. Stellar + * trustline fields) is applied in {@link SnapDataSource.fetch} when clients + * call `getAssets` after trustline-related events such as + * `AccountsController:accountAssetListUpdated`. + * * @param payload - The balance update payload from AccountsController. */ #handleSnapBalancesUpdated( @@ -303,11 +335,13 @@ export class SnapDataSource extends AbstractDataSource< } // Only report if we have snap-related updates - if (assetsBalance) { - const response: DataResponse = { assetsBalance, updateMode: 'merge' }; - for (const subscription of this.activeSubscriptions.values()) { - subscription.onAssetsUpdate(response)?.catch(console.error); - } + if (!assetsBalance) { + return; + } + + const response: DataResponse = { assetsBalance, updateMode: 'merge' }; + for (const subscription of this.activeSubscriptions.values()) { + subscription.onAssetsUpdate(response)?.catch(console.error); } } @@ -321,6 +355,150 @@ export class SnapDataSource extends AbstractDataSource< return this.state.activeChains.includes(chainId); } + async #enrichAccountAssetInfo( + assetsBalance: NonNullable, + ): Promise { + for (const [accountId, accountAssets] of Object.entries(assetsBalance)) { + const assetsByChain = new Map(); + + for (const assetId of Object.keys(accountAssets) as Caip19AssetId[]) { + let chainId: ChainId; + try { + chainId = extractChainFromAssetId(assetId); + } catch { + continue; + } + if (!isAccountAssetInfoEnrichmentAvailable(chainId)) { + continue; + } + + const assetIds = assetsByChain.get(chainId) ?? []; + assetIds.push(assetId); + assetsByChain.set(chainId, assetIds); + } + + for (const [chainId, assetIds] of assetsByChain) { + const snapId = this.getSnapIdForChain(chainId); + if (!snapId) { + continue; + } + + for ( + let i = 0; + i < assetIds.length; + i += ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE + ) { + const batch = assetIds.slice( + i, + i + ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE, + ); + const accountAssetInfo = await this.#fetchAccountAssetInfoFromSnap({ + accountId, + snapId, + chainId, + assets: batch, + }); + + if (!accountAssetInfo) { + break; + } + + for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { + const row = accountAssets[assetId as Caip19AssetId]; + if (!row) { + continue; + } + + accountAssets[assetId as Caip19AssetId] = { + ...row, + accountAssetInfo: assetInfo, + }; + } + } + } + } + } + + async #fetchAccountAssetInfoFromSnap({ + accountId, + snapId, + chainId, + assets, + }: { + accountId: string; + snapId: SnapId; + chainId: ChainId; + assets: Caip19AssetId[]; + }): Promise { + if (assets.length === 0) { + return undefined; + } + + try { + const request = { + snapId, + origin: 'metamask', + handler: HandlerType.OnClientRequest, + request: { + jsonrpc: '2.0', + method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + params: { + accountId, + scope: chainId, + assets, + }, + }, + }; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const snapRequest = (this.#messenger as any).call( + 'SnapController:handleRequest', + request, + ) as Promise; + + const timedOut = Symbol('timedOut'); + const result = await Promise.race([ + snapRequest, + new Promise((resolve) => + setTimeout(() => resolve(timedOut), ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS), + ), + ]); + + if (result === timedOut) { + snapRequest + .then(() => + log('Snap account asset info resolved after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + }), + ) + .catch((error) => + log('Snap account asset info failed after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }), + ); + return undefined; + } + + return result as GetAccountAssetInfoResponse | undefined; + } catch (error) { + log('Failed to enrich snap account asset info', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }); + return undefined; + } + } + // ============================================================================ // SNAP DISCOVERY (Dynamic via PermissionController) // ============================================================================ @@ -475,37 +653,57 @@ export class SnapDataSource extends AbstractDataSource< const client = this.#getKeyringClient(snapId); // Step 1: Get the list of assets for this account - const accountAssets = await client.listAccountAssets(accountId); + const listedAssets = (await client.listAccountAssets(accountId)) ?? []; + + // Include custom assets on enrichment chains (e.g. Stellar) so trustline + // fields refresh even when a deactivated asset drops off listAccountAssets. + const assetsToFetch = new Set( + listedAssets as CaipAssetType[], + ); + if (request.customAssets) { + for (const assetId of request.customAssets) { + try { + const assetChainId = extractChainFromAssetId(assetId); + if ( + request.chainIds.includes(assetChainId) && + isAccountAssetInfoEnrichmentAvailable(assetChainId) + ) { + assetsToFetch.add(assetId as CaipAssetType); + } + } catch { + // Ignore malformed custom asset ids. + } + } + } - // If no assets, skip to next account - if (!accountAssets || accountAssets.length === 0) { + const assetList = [...assetsToFetch]; + if (assetList.length === 0) { continue; } // Step 2: Get balances for those specific assets const balances: Record = - await client.getAccountBalances( - accountId, - accountAssets as CaipAssetType[], - ); + await client.getAccountBalances(accountId, assetList); // Transform keyring response to DataResponse format - if (balances && typeof balances === 'object' && results.assetsBalance) { - for (const [assetId, balance] of Object.entries(balances)) { - results.assetsBalance[accountId] ??= {}; - const accountBalances = results.assetsBalance[accountId]; - if (accountBalances) { - (accountBalances as Record)[assetId] = { - amount: balance.amount, - }; - } - } + results.assetsBalance ??= {}; + const accountBalances = results.assetsBalance[accountId] ?? {}; + for (const assetId of assetList) { + const balance = balances?.[assetId]; + (accountBalances as Record)[assetId] = { + amount: balance?.amount ?? '0', + }; } + results.assetsBalance[accountId] = accountBalances; } catch { // Expected when account doesn't belong to this snap } } + if (results.assetsBalance) { + await this.#enrichAccountAssetInfo(results.assetsBalance); + } + return results; } @@ -678,6 +876,17 @@ export class SnapDataSource extends AbstractDataSource< } } + /** + * Returns the snap id responsible for a given chain, if any. + * + * @param chainId - CAIP-2 chain id. + * @returns Snap id when a keyring snap supports the chain. + */ + getSnapIdForChain(chainId: ChainId): SnapId | undefined { + const snapId = this.state.chainToSnap[chainId]; + return snapId ? (snapId as SnapId) : undefined; + } + // ============================================================================ // KEYRING CLIENT // ============================================================================ diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index 0fb743e853..0fc27c1a8d 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -269,12 +269,31 @@ export type AssetPrice = FungibleAssetPrice | NFTAssetPrice; // BALANCE TYPES (vary by asset type) // ============================================================================ +/** + * Optional per-asset fields from snap account-asset enrichment. + * Stellar classic assets use trustline fields; native XLM may include baseReserve. + */ +export type AccountAssetInfo = { + limit?: string; + authorized?: boolean; + sponsored?: boolean; + baseReserve?: string; +}; + +/** Per-asset enrichment keyed by CAIP-19 asset id from snap getAccountAssetInfo. */ +export type GetAccountAssetInfoResponse = Record< + Caip19AssetId, + AccountAssetInfo +>; + /** * Balance data for fungible tokens (native, ERC20, SPL). */ export type FungibleAssetBalance = { /** Raw balance amount as string (e.g., "1000000000" for 1000 USDC) */ amount: string; + /** Optional per-asset info from snap account-asset enrichment (e.g. Stellar trustlines). */ + accountAssetInfo?: AccountAssetInfo; }; /** From e20f4b2724e28f24fed48519ca5e9512153ab66f Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Wed, 8 Jul 2026 02:24:34 +0200 Subject: [PATCH 2/9] chore: separate enrich logic in helper --- .../src/data-sources/SnapDataSource.test.ts | 121 --------- .../src/data-sources/SnapDataSource.ts | 213 ++------------- ...snap-account-asset-info-enrichment.test.ts | 205 +++++++++++++++ .../snap-account-asset-info-enrichment.ts | 246 ++++++++++++++++++ 4 files changed, 470 insertions(+), 315 deletions(-) create mode 100644 packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts create mode 100644 packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index 791ee402ff..64600bdcaa 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -674,127 +674,6 @@ describe('SnapDataSource', () => { cleanup(); }); - it('fetch returns balance without accountAssetInfo when snap hangs past timeout', async () => { - jest.useFakeTimers(); - - const stellarAccount = createMockAccount({ - scopes: [STELLAR_PUBNET], - metadata: { - name: 'Stellar Account', - importTime: Date.now(), - keyring: { type: 'Snap Keyring' }, - snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, - }, - }); - - const { controller, mockHandleRequest, cleanup } = setupController({ - installedSnaps: { - [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, - }, - accountAssets: [MOCK_STELLAR_USDC_ASSET], - balances: { [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' } }, - }); - - // Drain all microtasks from subscribe's initial fetch before replacing the mock - await jest.advanceTimersByTimeAsync(20_000); - - // Now configure getAccountAssetInfo to never resolve - mockHandleRequest.mockImplementation((params) => { - const { request } = params as { request: { method: string } }; - if (request?.method === 'keyring_listAccountAssets') { - return Promise.resolve([MOCK_STELLAR_USDC_ASSET]); - } - if (request?.method === 'keyring_getAccountBalances') { - return Promise.resolve({ - [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, - }); - } - if (request?.method === 'getAccountAssetInfo') { - return new Promise(() => undefined); // never resolves - } - return Promise.resolve(null); - }); - - const fetchPromise = controller.fetch( - createDataRequest({ chainIds: [STELLAR_PUBNET], accounts: [stellarAccount] }), - ); - - // Advance past the enrichment timeout - await jest.advanceTimersByTimeAsync(20_000); - const response = await fetchPromise; - - expect( - response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], - ).toStrictEqual({ amount: '25' }); - expect( - (response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET] as Record)?.accountAssetInfo, - ).toBeUndefined(); - - cleanup(); - jest.useRealTimers(); - }); - - it('fetch stops queuing batches after first batch times out', async () => { - jest.useFakeTimers(); - - const MOCK_STELLAR_ASSET_2 = - 'stellar:pubnet/asset:USDT-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; - const allAssets = [MOCK_STELLAR_USDC_ASSET, MOCK_STELLAR_ASSET_2]; - const balances = Object.fromEntries( - allAssets.map((a) => [a, { amount: '10', unit: 'X' }]), - ) as Record; - - const stellarAccount = createMockAccount({ - scopes: [STELLAR_PUBNET], - metadata: { - name: 'Stellar Account', - importTime: Date.now(), - keyring: { type: 'Snap Keyring' }, - snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, - }, - }); - - const { controller, mockHandleRequest, cleanup } = setupController({ - installedSnaps: { - [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, - }, - accountAssets: allAssets, - balances, - }); - - // Drain timers from subscribe's initial fetch - await jest.advanceTimersByTimeAsync(20_000); - - let getAccountAssetInfoCallCount = 0; - mockHandleRequest.mockImplementation((params) => { - const { request } = params as { request: { method: string } }; - if (request?.method === 'keyring_listAccountAssets') { - return Promise.resolve(allAssets); - } - if (request?.method === 'keyring_getAccountBalances') { - return Promise.resolve(balances); - } - if (request?.method === 'getAccountAssetInfo') { - getAccountAssetInfoCallCount += 1; - return new Promise(() => undefined); // always hangs - } - return Promise.resolve(null); - }); - - const fetchPromise = controller.fetch( - createDataRequest({ chainIds: [STELLAR_PUBNET], accounts: [stellarAccount] }), - ); - - await jest.advanceTimersByTimeAsync(20_000); - await fetchPromise; - - // Only 1 getAccountAssetInfo call: batch loop must break on timeout, not continue - expect(getAccountAssetInfoCallCount).toBe(1); - - cleanup(); - jest.useRealTimers(); - }); - it('fetch handles empty account assets gracefully', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ installedSnaps: { diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index e0db593540..cbcaa34d0e 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -25,10 +25,13 @@ import type { Caip19AssetId, DataRequest, DataResponse, - GetAccountAssetInfoResponse, Middleware, } from '../types'; import { AbstractDataSource } from './AbstractDataSource'; +import { + enrichAccountAssetInfo, + getAssetsToFetchWithEligibleCustomAssets, +} from './snap-account-asset-info-enrichment'; import type { DataSourceState, SubscriptionRequest, @@ -75,28 +78,6 @@ export const KEYRING_PERMISSION = 'endowment:keyring'; /** The permission name for snap assets endowment (contains chainIds) */ export const ASSETS_PERMISSION = 'endowment:assets'; -/** Snap clientRequest method for per-(account, asset) enrichment data. */ -const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; - -/** - * Max assets per snap getAccountAssetInfo request. Large batches can terminate - * the Stellar wallet snap on mobile when many trustlines are fetched at once. - */ -const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 3; - -/** Per-batch snap client request timeout (ms). Hung requests must not block apply. */ -const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; - -/** - * Chains whose wallet snap implements getAccountAssetInfo. - */ -const ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN: Partial< - Record -> = { - 'stellar:pubnet': true, - 'stellar:testnet': true, -}; - // ============================================================================ // PERMISSION UTILITIES // ============================================================================ @@ -138,10 +119,6 @@ export function extractChainFromAssetId(assetId: string): ChainId { return parsed.chainId; } -function isAccountAssetInfoEnrichmentAvailable(chainId: ChainId): boolean { - return ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN[chainId] === true; -} - // ============================================================================ // STATE // ============================================================================ @@ -355,150 +332,6 @@ export class SnapDataSource extends AbstractDataSource< return this.state.activeChains.includes(chainId); } - async #enrichAccountAssetInfo( - assetsBalance: NonNullable, - ): Promise { - for (const [accountId, accountAssets] of Object.entries(assetsBalance)) { - const assetsByChain = new Map(); - - for (const assetId of Object.keys(accountAssets) as Caip19AssetId[]) { - let chainId: ChainId; - try { - chainId = extractChainFromAssetId(assetId); - } catch { - continue; - } - if (!isAccountAssetInfoEnrichmentAvailable(chainId)) { - continue; - } - - const assetIds = assetsByChain.get(chainId) ?? []; - assetIds.push(assetId); - assetsByChain.set(chainId, assetIds); - } - - for (const [chainId, assetIds] of assetsByChain) { - const snapId = this.getSnapIdForChain(chainId); - if (!snapId) { - continue; - } - - for ( - let i = 0; - i < assetIds.length; - i += ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE - ) { - const batch = assetIds.slice( - i, - i + ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE, - ); - const accountAssetInfo = await this.#fetchAccountAssetInfoFromSnap({ - accountId, - snapId, - chainId, - assets: batch, - }); - - if (!accountAssetInfo) { - break; - } - - for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { - const row = accountAssets[assetId as Caip19AssetId]; - if (!row) { - continue; - } - - accountAssets[assetId as Caip19AssetId] = { - ...row, - accountAssetInfo: assetInfo, - }; - } - } - } - } - } - - async #fetchAccountAssetInfoFromSnap({ - accountId, - snapId, - chainId, - assets, - }: { - accountId: string; - snapId: SnapId; - chainId: ChainId; - assets: Caip19AssetId[]; - }): Promise { - if (assets.length === 0) { - return undefined; - } - - try { - const request = { - snapId, - origin: 'metamask', - handler: HandlerType.OnClientRequest, - request: { - jsonrpc: '2.0', - method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, - params: { - accountId, - scope: chainId, - assets, - }, - }, - }; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const snapRequest = (this.#messenger as any).call( - 'SnapController:handleRequest', - request, - ) as Promise; - - const timedOut = Symbol('timedOut'); - const result = await Promise.race([ - snapRequest, - new Promise((resolve) => - setTimeout(() => resolve(timedOut), ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS), - ), - ]); - - if (result === timedOut) { - snapRequest - .then(() => - log('Snap account asset info resolved after timeout', { - accountId, - snapId, - chainId, - assetCount: assets.length, - }), - ) - .catch((error) => - log('Snap account asset info failed after timeout', { - accountId, - snapId, - chainId, - assetCount: assets.length, - error, - }), - ); - return undefined; - } - - return result as GetAccountAssetInfoResponse | undefined; - } catch (error) { - log('Failed to enrich snap account asset info', { - accountId, - snapId, - chainId, - assetCount: assets.length, - error, - }); - return undefined; - } - } - // ============================================================================ // SNAP DISCOVERY (Dynamic via PermissionController) // ============================================================================ @@ -655,28 +488,12 @@ export class SnapDataSource extends AbstractDataSource< // Step 1: Get the list of assets for this account const listedAssets = (await client.listAccountAssets(accountId)) ?? []; - // Include custom assets on enrichment chains (e.g. Stellar) so trustline - // fields refresh even when a deactivated asset drops off listAccountAssets. - const assetsToFetch = new Set( - listedAssets as CaipAssetType[], - ); - if (request.customAssets) { - for (const assetId of request.customAssets) { - try { - const assetChainId = extractChainFromAssetId(assetId); - if ( - request.chainIds.includes(assetChainId) && - isAccountAssetInfoEnrichmentAvailable(assetChainId) - ) { - assetsToFetch.add(assetId as CaipAssetType); - } - } catch { - // Ignore malformed custom asset ids. - } - } - } - - const assetList = [...assetsToFetch]; + // TODO(STELLAR): Remove custom asset inclusion for enrichment once the Accounts API can refresh trustline/accountAssetInfo data for custom assets directly. + const assetList = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: listedAssets as CaipAssetType[], + customAssets: request.customAssets, + requestedChainIds: request.chainIds, + }); if (assetList.length === 0) { continue; } @@ -701,7 +518,15 @@ export class SnapDataSource extends AbstractDataSource< } if (results.assetsBalance) { - await this.#enrichAccountAssetInfo(results.assetsBalance); + // TODO(STELLAR): Remove this Snap-side accountAssetInfo enrichment path once the Accounts API returns account-asset enrichment directly. + await enrichAccountAssetInfo({ + assetsBalance: results.assetsBalance, + getSnapIdForChain: (chainId) => this.getSnapIdForChain(chainId), + callSnapRequest: (request) => + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (this.#messenger as any).call('SnapController:handleRequest', request), + log, + }); } return results; diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts new file mode 100644 index 0000000000..39c9decc0f --- /dev/null +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -0,0 +1,205 @@ +import type { CaipAssetType } from '@metamask/keyring-api'; +import type { SnapId } from '@metamask/snaps-sdk'; +import { HandlerType } from '@metamask/snaps-utils'; + +import type { ChainId, Caip19AssetId, DataResponse } from '../types'; +import { + GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + enrichAccountAssetInfo, + getAssetsToFetchWithEligibleCustomAssets, + isAccountAssetInfoEnrichmentAvailable, +} from './snap-account-asset-info-enrichment'; + +const STELLAR_PUBNET = 'stellar:pubnet' as ChainId; +const STELLAR_TESTNET = 'stellar:testnet' as ChainId; +const SOLANA_MAINNET = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; +const STELLAR_SNAP_ID = 'npm:@metamask/stellar-wallet-snap' as SnapId; + +const MOCK_STELLAR_USDC_ASSET = + 'stellar:pubnet/asset:USDC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; +const MOCK_STELLAR_AUDD_ASSET = + 'stellar:pubnet/asset:AUDD-GDC7X2MXTYSAKUUGAIQ7J7RPEIM7GXSAIWFYWWH4GLNFECQVJJLB2EEU' as Caip19AssetId; +const MOCK_STELLAR_ASSET_2 = + 'stellar:pubnet/asset:USDT-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; +const MOCK_SOL_ASSET = + 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501' as Caip19AssetId; + +describe('snap-account-asset-info-enrichment', () => { + describe('isAccountAssetInfoEnrichmentAvailable', () => { + it('returns true for Stellar chains', () => { + expect(isAccountAssetInfoEnrichmentAvailable(STELLAR_PUBNET)).toBe(true); + expect(isAccountAssetInfoEnrichmentAvailable(STELLAR_TESTNET)).toBe(true); + }); + + it('returns false for non-enrichment chains', () => { + expect(isAccountAssetInfoEnrichmentAvailable(SOLANA_MAINNET)).toBe(false); + }); + }); + + describe('getAssetsToFetchWithEligibleCustomAssets', () => { + it('includes eligible custom Stellar assets missing from listAccountAssets', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], + customAssets: [MOCK_STELLAR_AUDD_ASSET], + requestedChainIds: [STELLAR_PUBNET], + }); + + expect(result).toStrictEqual([ + MOCK_STELLAR_USDC_ASSET, + MOCK_STELLAR_AUDD_ASSET, + ]); + }); + + it('ignores custom assets on chains not in the request', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], + customAssets: [MOCK_STELLAR_AUDD_ASSET], + requestedChainIds: [SOLANA_MAINNET], + }); + + expect(result).toStrictEqual([MOCK_STELLAR_USDC_ASSET]); + }); + + it('ignores custom assets on enrichment-unavailable chains', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_SOL_ASSET as CaipAssetType], + customAssets: [MOCK_SOL_ASSET], + requestedChainIds: [SOLANA_MAINNET], + }); + + expect(result).toStrictEqual([MOCK_SOL_ASSET]); + }); + + it('ignores malformed custom asset ids', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], + customAssets: ['not-a-caip-asset' as Caip19AssetId], + requestedChainIds: [STELLAR_PUBNET], + }); + + expect(result).toStrictEqual([MOCK_STELLAR_USDC_ASSET]); + }); + }); + + describe('enrichAccountAssetInfo', () => { + const accountId = 'mock-account-id'; + + function createAssetsBalance(): NonNullable { + return { + [accountId]: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25' }, + }, + }; + } + + it('enriches Stellar assets with accountAssetInfo', async () => { + const assetsBalance = createAssetsBalance(); + const callSnapRequest = jest.fn().mockResolvedValue({ + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + sponsored: false, + }, + }); + + await enrichAccountAssetInfo({ + assetsBalance, + getSnapIdForChain: () => STELLAR_SNAP_ID, + callSnapRequest, + }); + + expect(callSnapRequest).toHaveBeenCalledWith({ + snapId: STELLAR_SNAP_ID, + origin: 'metamask', + handler: HandlerType.OnClientRequest, + request: { + jsonrpc: '2.0', + method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + params: { + accountId, + scope: STELLAR_PUBNET, + assets: [MOCK_STELLAR_USDC_ASSET], + }, + }, + }); + expect(assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '1000', + authorized: true, + sponsored: false, + }, + }); + }); + + it('returns balances without accountAssetInfo when enrichment hangs past timeout', async () => { + jest.useFakeTimers(); + + const assetsBalance = createAssetsBalance(); + const callSnapRequest = jest.fn( + () => new Promise(() => undefined), + ); + + const enrichPromise = enrichAccountAssetInfo({ + assetsBalance, + getSnapIdForChain: () => STELLAR_SNAP_ID, + callSnapRequest, + }); + + await jest.advanceTimersByTimeAsync(20_000); + await enrichPromise; + + expect(assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + }); + expect( + ( + assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET] as Record< + string, + unknown + > + )?.accountAssetInfo, + ).toBeUndefined(); + + jest.useRealTimers(); + }); + + it('stops after the first enrichment batch timeout', async () => { + jest.useFakeTimers(); + + const MOCK_STELLAR_ASSET_3 = + 'stellar:pubnet/asset:EURC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; + const MOCK_STELLAR_ASSET_4 = + 'stellar:pubnet/asset:BTC-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; + + const assetsBalance: NonNullable = { + [accountId]: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '10' }, + [MOCK_STELLAR_ASSET_2]: { amount: '10' }, + [MOCK_STELLAR_ASSET_3]: { amount: '10' }, + [MOCK_STELLAR_ASSET_4]: { amount: '10' }, + }, + }; + + let callCount = 0; + const callSnapRequest = jest.fn(() => { + callCount += 1; + return new Promise(() => undefined); + }); + + const enrichPromise = enrichAccountAssetInfo({ + assetsBalance, + getSnapIdForChain: () => STELLAR_SNAP_ID, + callSnapRequest, + }); + + await jest.advanceTimersByTimeAsync(20_000); + await enrichPromise; + + // Batch size is 3; a second batch would be queued without the timeout break. + expect(callCount).toBe(1); + + jest.useRealTimers(); + }); + }); +}); diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts new file mode 100644 index 0000000000..1fba4e887d --- /dev/null +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -0,0 +1,246 @@ +// TODO(STELLAR): This helper is a temporary bridge for Snap-provided accountAssetInfo. +// Remove it once the Accounts API supports account-asset enrichment directly. + +import type { CaipAssetType } from '@metamask/keyring-api'; +import type { SnapId } from '@metamask/snaps-sdk'; +import { HandlerType } from '@metamask/snaps-utils'; +import { parseCaipAssetType } from '@metamask/utils'; + +import { projectLogger, createModuleLogger } from '../logger'; +import type { + AssetBalance, + ChainId, + Caip19AssetId, + DataResponse, + GetAccountAssetInfoResponse, +} from '../types'; + +const log = createModuleLogger( + projectLogger, + 'snap-account-asset-info-enrichment', +); + +export const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; + +/** + * Max assets per snap getAccountAssetInfo request. Large batches can terminate + * the Stellar wallet snap on mobile when many trustlines are fetched at once. + */ +const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 3; + +/** Per-batch snap client request timeout (ms). Hung requests must not block apply. */ +const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; + +// TODO(STELLAR): Replace this chain allowlist with Accounts API-backed enrichment support. +const ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN: Partial< + Record +> = { + 'stellar:pubnet': true, + 'stellar:testnet': true, +}; + +const ENRICHMENT_TIMEOUT = Symbol('enrichmentTimeout'); + +function extractChainFromAssetId(assetId: string): ChainId { + const parsed = parseCaipAssetType(assetId as CaipAssetType); + return parsed.chainId; +} + +export function isAccountAssetInfoEnrichmentAvailable(chainId: ChainId): boolean { + return ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN[chainId] === true; +} + +export function getAssetsToFetchWithEligibleCustomAssets(params: { + listedAssets: CaipAssetType[]; + customAssets?: Caip19AssetId[]; + requestedChainIds: ChainId[]; +}): CaipAssetType[] { + const { listedAssets, customAssets, requestedChainIds } = params; + const assetsToFetch = new Set(listedAssets); + + if (customAssets) { + for (const assetId of customAssets) { + try { + const assetChainId = extractChainFromAssetId(assetId); + if ( + requestedChainIds.includes(assetChainId) && + isAccountAssetInfoEnrichmentAvailable(assetChainId) + ) { + assetsToFetch.add(assetId as CaipAssetType); + } + } catch { + // Ignore malformed custom asset ids. + } + } + } + + return [...assetsToFetch]; +} + +type SnapAccountAssetInfoRequest = { + snapId: SnapId; + origin: string; + handler: HandlerType; + request: { + jsonrpc: '2.0'; + method: string; + params: { + accountId: string; + scope: ChainId; + assets: Caip19AssetId[]; + }; + }; +}; + +async function fetchAccountAssetInfoFromSnap({ + accountId, + snapId, + chainId, + assets, + callSnapRequest, + moduleLog = log, +}: { + accountId: string; + snapId: SnapId; + chainId: ChainId; + assets: Caip19AssetId[]; + callSnapRequest: (request: SnapAccountAssetInfoRequest) => Promise; + moduleLog?: typeof log; +}): Promise { + if (assets.length === 0) { + return undefined; + } + + try { + const request: SnapAccountAssetInfoRequest = { + snapId, + origin: 'metamask', + handler: HandlerType.OnClientRequest, + request: { + jsonrpc: '2.0', + method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + params: { + accountId, + scope: chainId, + assets, + }, + }, + }; + + const snapRequest = callSnapRequest(request); + const result = await Promise.race([ + snapRequest, + new Promise((resolve) => + setTimeout(() => resolve(ENRICHMENT_TIMEOUT), ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS), + ), + ]); + + if (result === ENRICHMENT_TIMEOUT) { + snapRequest + .then(() => + moduleLog('Snap account asset info resolved after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + }), + ) + .catch((error) => + moduleLog('Snap account asset info failed after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }), + ); + return ENRICHMENT_TIMEOUT; + } + + return result as GetAccountAssetInfoResponse | undefined; + } catch (error) { + moduleLog('Failed to enrich snap account asset info', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }); + return ENRICHMENT_TIMEOUT; + } +} + +export async function enrichAccountAssetInfo(params: { + assetsBalance: NonNullable; + getSnapIdForChain: (chainId: ChainId) => SnapId | undefined; + callSnapRequest: (request: SnapAccountAssetInfoRequest) => Promise; + log?: typeof log; +}): Promise { + const { assetsBalance, getSnapIdForChain, callSnapRequest, log: moduleLog } = + params; + + for (const [accountId, accountAssets] of Object.entries(assetsBalance)) { + const assetsByChain = new Map(); + + for (const assetId of Object.keys(accountAssets) as Caip19AssetId[]) { + let chainId: ChainId; + try { + chainId = extractChainFromAssetId(assetId); + } catch { + continue; + } + if (!isAccountAssetInfoEnrichmentAvailable(chainId)) { + continue; + } + + const assetIds = assetsByChain.get(chainId) ?? []; + assetIds.push(assetId); + assetsByChain.set(chainId, assetIds); + } + + for (const [chainId, assetIds] of assetsByChain) { + const snapId = getSnapIdForChain(chainId); + if (!snapId) { + continue; + } + + for ( + let i = 0; + i < assetIds.length; + i += ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE + ) { + const batch = assetIds.slice( + i, + i + ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE, + ); + const accountAssetInfo = await fetchAccountAssetInfoFromSnap({ + accountId, + snapId, + chainId, + assets: batch, + callSnapRequest, + moduleLog, + }); + + if ( + accountAssetInfo === ENRICHMENT_TIMEOUT || + accountAssetInfo === undefined + ) { + break; + } + + for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { + const row = accountAssets[assetId as Caip19AssetId]; + if (!row) { + continue; + } + + accountAssets[assetId as Caip19AssetId] = { + ...row, + accountAssetInfo: assetInfo, + } satisfies AssetBalance; + } + } + } + } +} From dac217f843e6a302a0ddedee58716cece42c9f27 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Wed, 8 Jul 2026 02:29:24 +0200 Subject: [PATCH 3/9] fix: skip snap enrichment when unavailable --- .../src/data-sources/SnapDataSource.test.ts | 51 +++++++++++++++++ .../src/data-sources/SnapDataSource.ts | 9 ++- ...snap-account-asset-info-enrichment.test.ts | 57 +++++++++++++++++++ .../snap-account-asset-info-enrichment.ts | 27 +++++++++ 4 files changed, 143 insertions(+), 1 deletion(-) diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index 64600bdcaa..aeabd8e2db 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -606,6 +606,57 @@ describe('SnapDataSource', () => { cleanup(); }); + it('fetch does not call getAccountAssetInfo when no snap is available for enrichment chain', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ + installedSnaps: { + [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, + }, + accountAssets: [MOCK_STELLAR_USDC_ASSET], + balances: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, + }, + accountAssetInfo: { + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + }, + }, + }); + await new Promise(process.nextTick); + + jest.spyOn(controller, 'getSnapIdForChain').mockReturnValue(undefined); + + const response = await controller.fetch( + createDataRequest({ + chainIds: [STELLAR_PUBNET], + accounts: [ + createMockAccount({ + scopes: [STELLAR_PUBNET], + metadata: { + name: 'Stellar Account', + importTime: Date.now(), + keyring: { type: 'Snap Keyring' }, + snap: { id: STELLAR_SNAP_ID, name: 'Stellar Snap' }, + }, + }), + ], + }), + ); + + expect(mockHandleRequest).not.toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + method: 'getAccountAssetInfo', + }), + }), + ); + expect( + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], + ).toStrictEqual({ amount: '25' }); + + cleanup(); + }); + it('fetch enriches custom Stellar assets missing from listAccountAssets', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ installedSnaps: { diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index cbcaa34d0e..90720b7cbf 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -31,6 +31,7 @@ import { AbstractDataSource } from './AbstractDataSource'; import { enrichAccountAssetInfo, getAssetsToFetchWithEligibleCustomAssets, + hasAccountAssetInfoEnrichmentCandidate, } from './snap-account-asset-info-enrichment'; import type { DataSourceState, @@ -517,7 +518,13 @@ export class SnapDataSource extends AbstractDataSource< } } - if (results.assetsBalance) { + if ( + results.assetsBalance && + hasAccountAssetInfoEnrichmentCandidate({ + assetsBalance: results.assetsBalance, + getSnapIdForChain: (chainId) => this.getSnapIdForChain(chainId), + }) + ) { // TODO(STELLAR): Remove this Snap-side accountAssetInfo enrichment path once the Accounts API returns account-asset enrichment directly. await enrichAccountAssetInfo({ assetsBalance: results.assetsBalance, diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts index 39c9decc0f..cca3620abe 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -7,6 +7,7 @@ import { GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, enrichAccountAssetInfo, getAssetsToFetchWithEligibleCustomAssets, + hasAccountAssetInfoEnrichmentCandidate, isAccountAssetInfoEnrichmentAvailable, } from './snap-account-asset-info-enrichment'; @@ -81,6 +82,62 @@ describe('snap-account-asset-info-enrichment', () => { }); }); + describe('hasAccountAssetInfoEnrichmentCandidate', () => { + const accountId = 'mock-account-id'; + + it('returns true for Stellar asset when Snap ID exists', () => { + const result = hasAccountAssetInfoEnrichmentCandidate({ + assetsBalance: { + [accountId]: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25' }, + }, + }, + getSnapIdForChain: () => STELLAR_SNAP_ID, + }); + + expect(result).toBe(true); + }); + + it('returns false for Stellar asset when no Snap ID exists', () => { + const result = hasAccountAssetInfoEnrichmentCandidate({ + assetsBalance: { + [accountId]: { + [MOCK_STELLAR_USDC_ASSET]: { amount: '25' }, + }, + }, + getSnapIdForChain: () => undefined, + }); + + expect(result).toBe(false); + }); + + it('returns false for non-Stellar asset', () => { + const result = hasAccountAssetInfoEnrichmentCandidate({ + assetsBalance: { + [accountId]: { + [MOCK_SOL_ASSET]: { amount: '100' }, + }, + }, + getSnapIdForChain: () => STELLAR_SNAP_ID, + }); + + expect(result).toBe(false); + }); + + it('ignores malformed asset IDs', () => { + const result = hasAccountAssetInfoEnrichmentCandidate({ + assetsBalance: { + [accountId]: { + ['not-a-caip-asset' as Caip19AssetId]: { amount: '1' }, + }, + }, + getSnapIdForChain: () => STELLAR_SNAP_ID, + }); + + expect(result).toBe(false); + }); + }); + describe('enrichAccountAssetInfo', () => { const accountId = 'mock-account-id'; diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts index 1fba4e887d..fa2ad45cd5 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -50,6 +50,33 @@ export function isAccountAssetInfoEnrichmentAvailable(chainId: ChainId): boolean return ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN[chainId] === true; } +export function hasAccountAssetInfoEnrichmentCandidate(params: { + assetsBalance: NonNullable; + getSnapIdForChain: (chainId: ChainId) => SnapId | undefined; +}): boolean { + const { assetsBalance, getSnapIdForChain } = params; + + for (const accountAssets of Object.values(assetsBalance)) { + for (const assetId of Object.keys(accountAssets) as Caip19AssetId[]) { + let chainId: ChainId; + try { + chainId = extractChainFromAssetId(assetId); + } catch { + continue; + } + + if ( + isAccountAssetInfoEnrichmentAvailable(chainId) && + getSnapIdForChain(chainId) + ) { + return true; + } + } + } + + return false; +} + export function getAssetsToFetchWithEligibleCustomAssets(params: { listedAssets: CaipAssetType[]; customAssets?: Caip19AssetId[]; From 31bb1f3e5aa1ea0bff1fd21ffcc8dd2e4a6f7af0 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Wed, 8 Jul 2026 14:46:52 +0200 Subject: [PATCH 4/9] chore: simplify --- packages/assets-controller/CHANGELOG.md | 1 + .../assets-controller/src/AssetsController.ts | 9 +- .../src/data-sources/SnapDataSource.test.ts | 54 +++++------- .../src/data-sources/SnapDataSource.ts | 83 ++++++++++--------- .../src/data-sources/index.ts | 1 + ...snap-account-asset-info-enrichment.test.ts | 49 ----------- .../snap-account-asset-info-enrichment.ts | 29 +------ packages/assets-controller/src/index.ts | 1 + 8 files changed, 79 insertions(+), 148 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index d437a721f9..0b6ceba901 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add optional `snapDataSourceConfig.assetEnrichmentEnabled` getter to `AssetsControllerOptions` so clients can enable or disable snap `getAccountAssetInfo` enrichment at runtime (defaults to disabled) - Add optional snap account-asset enrichment to fungible balances on `SnapDataSource.fetch` so Stellar trustline fields (`accountAssetInfo`) are included when clients call `getAssets`; balance push events (`AccountsController:accountBalancesUpdated`) pass amounts through without enrichment — clients should call `getAssets` after trustline-related snap events (e.g. `AccountsController:accountAssetListUpdated`) ### Changed diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 8e558d7c2c..eece474d7b 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -84,7 +84,10 @@ import type { PriceDataSourceConfig } from './data-sources/PriceDataSource'; import { PriceDataSource } from './data-sources/PriceDataSource'; import type { RpcDataSourceConfig } from './data-sources/RpcDataSource'; import { RpcDataSource } from './data-sources/RpcDataSource'; -import type { AccountsControllerAccountBalancesUpdatedEvent } from './data-sources/SnapDataSource'; +import type { + AccountsControllerAccountBalancesUpdatedEvent, + SnapDataSourceConfig, +} from './data-sources/SnapDataSource'; import { SnapDataSource } from './data-sources/SnapDataSource'; import type { StakedBalanceDataSourceConfig } from './data-sources/StakedBalanceDataSource'; import { StakedBalanceDataSource } from './data-sources/StakedBalanceDataSource'; @@ -428,6 +431,8 @@ export type AssetsControllerOptions = { priceDataSourceConfig?: PriceDataSourceConfig; /** Optional configuration for StakedBalanceDataSource. */ stakedBalanceDataSourceConfig?: StakedBalanceDataSourceConfig; + /** Optional configuration for SnapDataSource. */ + snapDataSourceConfig?: SnapDataSourceConfig; /** * Function returning whether onboarding is complete. When false, * RPC and staked balance data sources skip fetch and subscribe @@ -855,6 +860,7 @@ export class AssetsController extends BaseController< accountsApiDataSourceConfig, priceDataSourceConfig, stakedBalanceDataSourceConfig, + snapDataSourceConfig, isOnboarded, }: AssetsControllerOptions) { super({ @@ -907,6 +913,7 @@ export class AssetsController extends BaseController< this.#snapDataSource = new SnapDataSource({ messenger: this.messenger, onActiveChainsUpdated: this.#onActiveChainsUpdated, + ...snapDataSourceConfig, }); this.#rpcDataSource = new RpcDataSource({ messenger: this.messenger, diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index aeabd8e2db..ef9710afe9 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -186,6 +186,7 @@ function setupController( balances?: Record; accountAssetInfo?: Record; configuredNetworks?: ChainId[]; + assetEnrichmentEnabled?: () => boolean; } = {}, ): SetupResult { const { @@ -193,6 +194,7 @@ function setupController( accountAssets = [], balances = {}, accountAssetInfo = {}, + assetEnrichmentEnabled, } = options; const rootMessenger = new Messenger({ @@ -269,6 +271,9 @@ function setupController( const controllerOptions: SnapDataSourceOptions = { messenger: controllerMessenger as unknown as AssetsControllerMessenger, onActiveChainsUpdated: activeChainsUpdateHandler, + ...(assetEnrichmentEnabled === undefined + ? {} + : { assetEnrichmentEnabled }), }; const controller = new SnapDataSource(controllerOptions); @@ -543,6 +548,7 @@ describe('SnapDataSource', () => { it('fetch enriches Stellar assets with account asset info', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ + assetEnrichmentEnabled: () => true, installedSnaps: { [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, }, @@ -606,8 +612,9 @@ describe('SnapDataSource', () => { cleanup(); }); - it('fetch does not call getAccountAssetInfo when no snap is available for enrichment chain', async () => { + it('fetch does not include customAssets in the snap balance request', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ + assetEnrichmentEnabled: () => true, installedSnaps: { [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, }, @@ -624,11 +631,10 @@ describe('SnapDataSource', () => { }); await new Promise(process.nextTick); - jest.spyOn(controller, 'getSnapIdForChain').mockReturnValue(undefined); - const response = await controller.fetch( createDataRequest({ chainIds: [STELLAR_PUBNET], + customAssets: [MOCK_STELLAR_AUDD_ASSET], accounts: [ createMockAccount({ scopes: [STELLAR_PUBNET], @@ -643,39 +649,39 @@ describe('SnapDataSource', () => { }), ); - expect(mockHandleRequest).not.toHaveBeenCalledWith( + expect(mockHandleRequest).toHaveBeenCalledWith( expect.objectContaining({ request: expect.objectContaining({ - method: 'getAccountAssetInfo', + method: 'keyring_getAccountBalances', + params: { + id: 'mock-account-id', + assets: [MOCK_STELLAR_USDC_ASSET], + }, }), }), ); expect( - response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], - ).toStrictEqual({ amount: '25' }); + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_AUDD_ASSET], + ).toBeUndefined(); cleanup(); }); - it('fetch enriches custom Stellar assets missing from listAccountAssets', async () => { + it('fetch does not call getAccountAssetInfo when assetEnrichmentEnabled is false', async () => { const { controller, mockHandleRequest, cleanup } = setupController({ + assetEnrichmentEnabled: () => false, installedSnaps: { [STELLAR_SNAP_ID]: { version: '1.0.0', chainIds: [STELLAR_PUBNET] }, }, accountAssets: [MOCK_STELLAR_USDC_ASSET], balances: { [MOCK_STELLAR_USDC_ASSET]: { amount: '25', unit: 'USDC' }, - [MOCK_STELLAR_AUDD_ASSET]: { amount: '0', unit: 'AUDD' }, }, accountAssetInfo: { [MOCK_STELLAR_USDC_ASSET]: { limit: '1000', authorized: true, }, - [MOCK_STELLAR_AUDD_ASSET]: { - limit: '0', - authorized: true, - }, }, }); await new Promise(process.nextTick); @@ -683,7 +689,6 @@ describe('SnapDataSource', () => { const response = await controller.fetch( createDataRequest({ chainIds: [STELLAR_PUBNET], - customAssets: [MOCK_STELLAR_AUDD_ASSET], accounts: [ createMockAccount({ scopes: [STELLAR_PUBNET], @@ -698,29 +703,16 @@ describe('SnapDataSource', () => { }), ); - expect(mockHandleRequest).toHaveBeenCalledWith( + expect(mockHandleRequest).not.toHaveBeenCalledWith( expect.objectContaining({ request: expect.objectContaining({ - method: 'keyring_getAccountBalances', - params: { - id: 'mock-account-id', - assets: expect.arrayContaining([ - MOCK_STELLAR_USDC_ASSET, - MOCK_STELLAR_AUDD_ASSET, - ]), - }, + method: 'getAccountAssetInfo', }), }), ); expect( - response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_AUDD_ASSET], - ).toStrictEqual({ - amount: '0', - accountAssetInfo: { - limit: '0', - authorized: true, - }, - }); + response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], + ).toStrictEqual({ amount: '25' }); cleanup(); }); diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index 90720b7cbf..d4e60d15dc 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -30,7 +30,6 @@ import type { import { AbstractDataSource } from './AbstractDataSource'; import { enrichAccountAssetInfo, - getAssetsToFetchWithEligibleCustomAssets, hasAccountAssetInfoEnrichmentCandidate, } from './snap-account-asset-info-enrichment'; import type { @@ -162,7 +161,15 @@ export type SnapDataSourceAllowedActions = // OPTIONS // ============================================================================ -export type SnapDataSourceOptions = { +export type SnapDataSourceConfig = { + /** + * When false, SnapDataSource skips `getAccountAssetInfo` enrichment. + * Evaluated at call time. Defaults to () => false. + */ + assetEnrichmentEnabled?: () => boolean; +}; + +export type SnapDataSourceOptions = SnapDataSourceConfig & { /** The AssetsController messenger (shared by all data sources). */ messenger: AssetsControllerMessenger; /** Called when this data source's active chains change. Pass dataSourceName so the controller knows the source. */ @@ -223,6 +230,8 @@ export class SnapDataSource extends AbstractDataSource< /** Cache of KeyringClient instances per snap ID to avoid re-instantiation */ readonly #keyringClientCache: Map = new Map(); + readonly #assetEnrichmentEnabled: () => boolean; + constructor(options: SnapDataSourceOptions) { super(SNAP_DATA_SOURCE_NAME, { ...defaultSnapState, @@ -231,6 +240,8 @@ export class SnapDataSource extends AbstractDataSource< this.#messenger = options.messenger; this.#onActiveChainsUpdated = options.onActiveChainsUpdated; + this.#assetEnrichmentEnabled = + options.assetEnrichmentEnabled ?? ((): boolean => false); // Bind handlers for cleanup in destroy() this.#handleSnapBalancesUpdatedBound = this.#handleSnapBalancesUpdated.bind( @@ -313,13 +324,11 @@ export class SnapDataSource extends AbstractDataSource< } // Only report if we have snap-related updates - if (!assetsBalance) { - return; - } - - const response: DataResponse = { assetsBalance, updateMode: 'merge' }; - for (const subscription of this.activeSubscriptions.values()) { - subscription.onAssetsUpdate(response)?.catch(console.error); + if (assetsBalance) { + const response: DataResponse = { assetsBalance, updateMode: 'merge' }; + for (const subscription of this.activeSubscriptions.values()) { + subscription.onAssetsUpdate(response)?.catch(console.error); + } } } @@ -487,48 +496,55 @@ export class SnapDataSource extends AbstractDataSource< const client = this.#getKeyringClient(snapId); // Step 1: Get the list of assets for this account - const listedAssets = (await client.listAccountAssets(accountId)) ?? []; + const accountAssets = await client.listAccountAssets(accountId); - // TODO(STELLAR): Remove custom asset inclusion for enrichment once the Accounts API can refresh trustline/accountAssetInfo data for custom assets directly. - const assetList = getAssetsToFetchWithEligibleCustomAssets({ - listedAssets: listedAssets as CaipAssetType[], - customAssets: request.customAssets, - requestedChainIds: request.chainIds, - }); - if (assetList.length === 0) { + // If no assets, skip to next account + if (!accountAssets || accountAssets.length === 0) { continue; } // Step 2: Get balances for those specific assets const balances: Record = - await client.getAccountBalances(accountId, assetList); + await client.getAccountBalances( + accountId, + accountAssets as CaipAssetType[], + ); // Transform keyring response to DataResponse format - results.assetsBalance ??= {}; - const accountBalances = results.assetsBalance[accountId] ?? {}; - for (const assetId of assetList) { - const balance = balances?.[assetId]; - (accountBalances as Record)[assetId] = { - amount: balance?.amount ?? '0', - }; + if (balances && typeof balances === 'object' && results.assetsBalance) { + for (const [assetId, balance] of Object.entries(balances)) { + results.assetsBalance[accountId] ??= {}; + const accountBalances = results.assetsBalance[accountId]; + if (accountBalances) { + (accountBalances as Record)[assetId] = { + amount: balance.amount, + }; + } + } } - results.assetsBalance[accountId] = accountBalances; } catch { // Expected when account doesn't belong to this snap } } + // Post-fetch enrichment stage: assetsBalance above already matches the + // standard (unenriched) balance shape. When the feature flag is enabled + // and there are eligible assets, apply accountAssetInfo enrichment once. + const getSnapIdForChain = (chainId: ChainId): SnapId | undefined => + this.state.chainToSnap[chainId] as SnapId | undefined; + if ( + this.#assetEnrichmentEnabled() && results.assetsBalance && hasAccountAssetInfoEnrichmentCandidate({ assetsBalance: results.assetsBalance, - getSnapIdForChain: (chainId) => this.getSnapIdForChain(chainId), + getSnapIdForChain, }) ) { // TODO(STELLAR): Remove this Snap-side accountAssetInfo enrichment path once the Accounts API returns account-asset enrichment directly. await enrichAccountAssetInfo({ assetsBalance: results.assetsBalance, - getSnapIdForChain: (chainId) => this.getSnapIdForChain(chainId), + getSnapIdForChain, callSnapRequest: (request) => // eslint-disable-next-line @typescript-eslint/no-explicit-any (this.#messenger as any).call('SnapController:handleRequest', request), @@ -708,17 +724,6 @@ export class SnapDataSource extends AbstractDataSource< } } - /** - * Returns the snap id responsible for a given chain, if any. - * - * @param chainId - CAIP-2 chain id. - * @returns Snap id when a keyring snap supports the chain. - */ - getSnapIdForChain(chainId: ChainId): SnapId | undefined { - const snapId = this.state.chainToSnap[chainId]; - return snapId ? (snapId as SnapId) : undefined; - } - // ============================================================================ // KEYRING CLIENT // ============================================================================ diff --git a/packages/assets-controller/src/data-sources/index.ts b/packages/assets-controller/src/data-sources/index.ts index 47cb7c5734..809b947ccf 100644 --- a/packages/assets-controller/src/data-sources/index.ts +++ b/packages/assets-controller/src/data-sources/index.ts @@ -57,6 +57,7 @@ export { // Types type SnapDataSourceState, type SnapDataSourceOptions, + type SnapDataSourceConfig, type SnapDataSourceAllowedActions, type SnapDataSourceAllowedEvents, } from './SnapDataSource'; diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts index cca3620abe..fc5da0164c 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -1,4 +1,3 @@ -import type { CaipAssetType } from '@metamask/keyring-api'; import type { SnapId } from '@metamask/snaps-sdk'; import { HandlerType } from '@metamask/snaps-utils'; @@ -6,7 +5,6 @@ import type { ChainId, Caip19AssetId, DataResponse } from '../types'; import { GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, enrichAccountAssetInfo, - getAssetsToFetchWithEligibleCustomAssets, hasAccountAssetInfoEnrichmentCandidate, isAccountAssetInfoEnrichmentAvailable, } from './snap-account-asset-info-enrichment'; @@ -18,8 +16,6 @@ const STELLAR_SNAP_ID = 'npm:@metamask/stellar-wallet-snap' as SnapId; const MOCK_STELLAR_USDC_ASSET = 'stellar:pubnet/asset:USDC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; -const MOCK_STELLAR_AUDD_ASSET = - 'stellar:pubnet/asset:AUDD-GDC7X2MXTYSAKUUGAIQ7J7RPEIM7GXSAIWFYWWH4GLNFECQVJJLB2EEU' as Caip19AssetId; const MOCK_STELLAR_ASSET_2 = 'stellar:pubnet/asset:USDT-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; const MOCK_SOL_ASSET = @@ -37,51 +33,6 @@ describe('snap-account-asset-info-enrichment', () => { }); }); - describe('getAssetsToFetchWithEligibleCustomAssets', () => { - it('includes eligible custom Stellar assets missing from listAccountAssets', () => { - const result = getAssetsToFetchWithEligibleCustomAssets({ - listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], - customAssets: [MOCK_STELLAR_AUDD_ASSET], - requestedChainIds: [STELLAR_PUBNET], - }); - - expect(result).toStrictEqual([ - MOCK_STELLAR_USDC_ASSET, - MOCK_STELLAR_AUDD_ASSET, - ]); - }); - - it('ignores custom assets on chains not in the request', () => { - const result = getAssetsToFetchWithEligibleCustomAssets({ - listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], - customAssets: [MOCK_STELLAR_AUDD_ASSET], - requestedChainIds: [SOLANA_MAINNET], - }); - - expect(result).toStrictEqual([MOCK_STELLAR_USDC_ASSET]); - }); - - it('ignores custom assets on enrichment-unavailable chains', () => { - const result = getAssetsToFetchWithEligibleCustomAssets({ - listedAssets: [MOCK_SOL_ASSET as CaipAssetType], - customAssets: [MOCK_SOL_ASSET], - requestedChainIds: [SOLANA_MAINNET], - }); - - expect(result).toStrictEqual([MOCK_SOL_ASSET]); - }); - - it('ignores malformed custom asset ids', () => { - const result = getAssetsToFetchWithEligibleCustomAssets({ - listedAssets: [MOCK_STELLAR_USDC_ASSET as CaipAssetType], - customAssets: ['not-a-caip-asset' as Caip19AssetId], - requestedChainIds: [STELLAR_PUBNET], - }); - - expect(result).toStrictEqual([MOCK_STELLAR_USDC_ASSET]); - }); - }); - describe('hasAccountAssetInfoEnrichmentCandidate', () => { const accountId = 'mock-account-id'; diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts index fa2ad45cd5..93e8f3fb4e 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -26,7 +26,7 @@ export const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; * Max assets per snap getAccountAssetInfo request. Large batches can terminate * the Stellar wallet snap on mobile when many trustlines are fetched at once. */ -const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 3; +const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 1; /** Per-batch snap client request timeout (ms). Hung requests must not block apply. */ const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; @@ -77,33 +77,6 @@ export function hasAccountAssetInfoEnrichmentCandidate(params: { return false; } -export function getAssetsToFetchWithEligibleCustomAssets(params: { - listedAssets: CaipAssetType[]; - customAssets?: Caip19AssetId[]; - requestedChainIds: ChainId[]; -}): CaipAssetType[] { - const { listedAssets, customAssets, requestedChainIds } = params; - const assetsToFetch = new Set(listedAssets); - - if (customAssets) { - for (const assetId of customAssets) { - try { - const assetChainId = extractChainFromAssetId(assetId); - if ( - requestedChainIds.includes(assetChainId) && - isAccountAssetInfoEnrichmentAvailable(assetChainId) - ) { - assetsToFetch.add(assetId as CaipAssetType); - } - } catch { - // Ignore malformed custom asset ids. - } - } - } - - return [...assetsToFetch]; -} - type SnapAccountAssetInfoRequest = { snapId: SnapId; origin: string; diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 1ebe4be876..06bc09c0f3 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -146,6 +146,7 @@ export { export type { SnapDataSourceState, SnapDataSourceOptions, + SnapDataSourceConfig, } from './data-sources'; // Enrichment data sources From f39d837c50e1330c8cfd4c1b90d860620ee32164 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Wed, 8 Jul 2026 15:01:56 +0200 Subject: [PATCH 5/9] chore: removed getSnapIdForChain --- .../assets-controller/src/data-sources/SnapDataSource.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index d4e60d15dc..7ec71e3f0a 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -530,21 +530,18 @@ export class SnapDataSource extends AbstractDataSource< // Post-fetch enrichment stage: assetsBalance above already matches the // standard (unenriched) balance shape. When the feature flag is enabled // and there are eligible assets, apply accountAssetInfo enrichment once. - const getSnapIdForChain = (chainId: ChainId): SnapId | undefined => - this.state.chainToSnap[chainId] as SnapId | undefined; - if ( this.#assetEnrichmentEnabled() && results.assetsBalance && hasAccountAssetInfoEnrichmentCandidate({ assetsBalance: results.assetsBalance, - getSnapIdForChain, + chainToSnap: this.state.chainToSnap, }) ) { // TODO(STELLAR): Remove this Snap-side accountAssetInfo enrichment path once the Accounts API returns account-asset enrichment directly. await enrichAccountAssetInfo({ assetsBalance: results.assetsBalance, - getSnapIdForChain, + chainToSnap: this.state.chainToSnap, callSnapRequest: (request) => // eslint-disable-next-line @typescript-eslint/no-explicit-any (this.#messenger as any).call('SnapController:handleRequest', request), From ac97644379ac037e79c596f614725591d8cdb8d7 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Thu, 9 Jul 2026 11:29:54 +0200 Subject: [PATCH 6/9] chore: rollback package.json --- packages/assets-controller/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/assets-controller/package.json b/packages/assets-controller/package.json index 834aa386ca..9fa1ae7755 100644 --- a/packages/assets-controller/package.json +++ b/packages/assets-controller/package.json @@ -1,6 +1,6 @@ { "name": "@metamask/assets-controller", - "version": "10.0.1-dev.8", + "version": "10.0.1", "description": "Tracks assets balances/prices and handles token detection across all digital assets", "keywords": [ "Ethereum", From 02607988218916c65b9ad1faf14a6ac625f3b029 Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Thu, 9 Jul 2026 13:14:19 +0200 Subject: [PATCH 7/9] chore: refactor snap asset metadata enrichment --- .../src/data-sources/SnapDataSource.ts | 37 +- ...snap-account-asset-info-enrichment.test.ts | 131 ++++--- .../snap-account-asset-info-enrichment.ts | 327 ++++++++++-------- 3 files changed, 278 insertions(+), 217 deletions(-) diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index 7ec71e3f0a..475a5e15fa 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -29,8 +29,8 @@ import type { } from '../types'; import { AbstractDataSource } from './AbstractDataSource'; import { - enrichAccountAssetInfo, hasAccountAssetInfoEnrichmentCandidate, + SnapAccountAssetInfoEnricher, } from './snap-account-asset-info-enrichment'; import type { DataSourceState, @@ -232,6 +232,8 @@ export class SnapDataSource extends AbstractDataSource< readonly #assetEnrichmentEnabled: () => boolean; + readonly #accountAssetInfoEnricher: SnapAccountAssetInfoEnricher; + constructor(options: SnapDataSourceOptions) { super(SNAP_DATA_SOURCE_NAME, { ...defaultSnapState, @@ -242,6 +244,15 @@ export class SnapDataSource extends AbstractDataSource< this.#onActiveChainsUpdated = options.onActiveChainsUpdated; this.#assetEnrichmentEnabled = options.assetEnrichmentEnabled ?? ((): boolean => false); + // TODO(STELLAR): Remove once the Accounts API returns account-asset metadata directly. + this.#accountAssetInfoEnricher = new SnapAccountAssetInfoEnricher({ + getSnapIdForChain: (chainId): SnapId | undefined => + this.state.chainToSnap[chainId] as SnapId | undefined, + callSnapRequest: (request): Promise => + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (this.#messenger as any).call('SnapController:handleRequest', request), + log, + }); // Bind handlers for cleanup in destroy() this.#handleSnapBalancesUpdatedBound = this.#handleSnapBalancesUpdated.bind( @@ -527,26 +538,24 @@ export class SnapDataSource extends AbstractDataSource< } } - // Post-fetch enrichment stage: assetsBalance above already matches the - // standard (unenriched) balance shape. When the feature flag is enabled - // and there are eligible assets, apply accountAssetInfo enrichment once. if ( this.#assetEnrichmentEnabled() && results.assetsBalance && hasAccountAssetInfoEnrichmentCandidate({ assetsBalance: results.assetsBalance, - chainToSnap: this.state.chainToSnap, + getSnapIdForChain: (chainId): SnapId | undefined => + this.state.chainToSnap[chainId] as SnapId | undefined, }) ) { - // TODO(STELLAR): Remove this Snap-side accountAssetInfo enrichment path once the Accounts API returns account-asset enrichment directly. - await enrichAccountAssetInfo({ - assetsBalance: results.assetsBalance, - chainToSnap: this.state.chainToSnap, - callSnapRequest: (request) => - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (this.#messenger as any).call('SnapController:handleRequest', request), - log, - }); + // TODO(STELLAR): Remove once the Accounts API returns account-asset metadata directly. + for (const [accountId, accountAssets] of Object.entries( + results.assetsBalance, + )) { + await this.#accountAssetInfoEnricher.enrichAccount({ + accountId, + assetsBalance: accountAssets, + }); + } } return results; diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts index fc5da0164c..59a907a6ee 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -4,9 +4,9 @@ import { HandlerType } from '@metamask/snaps-utils'; import type { ChainId, Caip19AssetId, DataResponse } from '../types'; import { GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, - enrichAccountAssetInfo, hasAccountAssetInfoEnrichmentCandidate, isAccountAssetInfoEnrichmentAvailable, + SnapAccountAssetInfoEnricher, } from './snap-account-asset-info-enrichment'; const STELLAR_PUBNET = 'stellar:pubnet' as ChainId; @@ -89,19 +89,28 @@ describe('snap-account-asset-info-enrichment', () => { }); }); - describe('enrichAccountAssetInfo', () => { + describe('SnapAccountAssetInfoEnricher', () => { const accountId = 'mock-account-id'; - function createAssetsBalance(): NonNullable { + function createAccountAssets(): NonNullable< + DataResponse['assetsBalance'] + >[string] { return { - [accountId]: { - [MOCK_STELLAR_USDC_ASSET]: { amount: '25' }, - }, + [MOCK_STELLAR_USDC_ASSET]: { amount: '25' }, }; } + function createEnricher( + callSnapRequest: jest.Mock, + ): SnapAccountAssetInfoEnricher { + return new SnapAccountAssetInfoEnricher({ + getSnapIdForChain: () => STELLAR_SNAP_ID, + callSnapRequest, + }); + } + it('enriches Stellar assets with accountAssetInfo', async () => { - const assetsBalance = createAssetsBalance(); + const assetsBalance = createAccountAssets(); const callSnapRequest = jest.fn().mockResolvedValue({ [MOCK_STELLAR_USDC_ASSET]: { limit: '1000', @@ -109,12 +118,9 @@ describe('snap-account-asset-info-enrichment', () => { sponsored: false, }, }); + const enricher = createEnricher(callSnapRequest); - await enrichAccountAssetInfo({ - assetsBalance, - getSnapIdForChain: () => STELLAR_SNAP_ID, - callSnapRequest, - }); + await enricher.enrichAccount({ accountId, assetsBalance }); expect(callSnapRequest).toHaveBeenCalledWith({ snapId: STELLAR_SNAP_ID, @@ -130,7 +136,7 @@ describe('snap-account-asset-info-enrichment', () => { }, }, }); - expect(assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', accountAssetInfo: { limit: '1000', @@ -140,74 +146,85 @@ describe('snap-account-asset-info-enrichment', () => { }); }); + it('requests all assets on a chain in a single getAccountAssetInfo call', async () => { + const assetsBalance: NonNullable< + DataResponse['assetsBalance'] + >[string] = { + [MOCK_STELLAR_USDC_ASSET]: { amount: '10' }, + [MOCK_STELLAR_ASSET_2]: { amount: '20' }, + }; + const callSnapRequest = jest.fn().mockResolvedValue({ + [MOCK_STELLAR_USDC_ASSET]: { limit: '1000', authorized: true }, + [MOCK_STELLAR_ASSET_2]: { limit: '500', authorized: false }, + }); + const enricher = createEnricher(callSnapRequest); + + await enricher.enrichAccount({ accountId, assetsBalance }); + + expect(callSnapRequest).toHaveBeenCalledTimes(1); + expect(callSnapRequest).toHaveBeenCalledWith( + expect.objectContaining({ + request: expect.objectContaining({ + params: expect.objectContaining({ + assets: expect.arrayContaining([ + MOCK_STELLAR_USDC_ASSET, + MOCK_STELLAR_ASSET_2, + ]), + }), + }), + }), + ); + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '10', + accountAssetInfo: { limit: '1000', authorized: true }, + }); + expect(assetsBalance[MOCK_STELLAR_ASSET_2]).toStrictEqual({ + amount: '20', + accountAssetInfo: { limit: '500', authorized: false }, + }); + }); + it('returns balances without accountAssetInfo when enrichment hangs past timeout', async () => { jest.useFakeTimers(); - const assetsBalance = createAssetsBalance(); - const callSnapRequest = jest.fn( - () => new Promise(() => undefined), - ); + const assetsBalance = createAccountAssets(); + const callSnapRequest = jest.fn(() => new Promise(() => undefined)); + const enricher = createEnricher(callSnapRequest); - const enrichPromise = enrichAccountAssetInfo({ + const enrichPromise = enricher.enrichAccount({ + accountId, assetsBalance, - getSnapIdForChain: () => STELLAR_SNAP_ID, - callSnapRequest, }); await jest.advanceTimersByTimeAsync(20_000); await enrichPromise; - expect(assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', }); expect( ( - assetsBalance[accountId]?.[MOCK_STELLAR_USDC_ASSET] as Record< - string, - unknown - > + assetsBalance[MOCK_STELLAR_USDC_ASSET] as Record )?.accountAssetInfo, ).toBeUndefined(); jest.useRealTimers(); }); - it('stops after the first enrichment batch timeout', async () => { - jest.useFakeTimers(); - - const MOCK_STELLAR_ASSET_3 = - 'stellar:pubnet/asset:EURC-GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN' as Caip19AssetId; - const MOCK_STELLAR_ASSET_4 = - 'stellar:pubnet/asset:BTC-GCQTGZQQ5G4PTM2GL7CDIFKUBIPEC52BROAQIAPW53XBRJVN6ZJVTG6' as Caip19AssetId; - - const assetsBalance: NonNullable = { - [accountId]: { - [MOCK_STELLAR_USDC_ASSET]: { amount: '10' }, - [MOCK_STELLAR_ASSET_2]: { amount: '10' }, - [MOCK_STELLAR_ASSET_3]: { amount: '10' }, - [MOCK_STELLAR_ASSET_4]: { amount: '10' }, - }, - }; - - let callCount = 0; - const callSnapRequest = jest.fn(() => { - callCount += 1; - return new Promise(() => undefined); - }); - - const enrichPromise = enrichAccountAssetInfo({ - assetsBalance, - getSnapIdForChain: () => STELLAR_SNAP_ID, + it('skips enrichment when no Snap ID is available for the chain', async () => { + const assetsBalance = createAccountAssets(); + const callSnapRequest = jest.fn(); + const enricher = new SnapAccountAssetInfoEnricher({ + getSnapIdForChain: (): SnapId | undefined => undefined, callSnapRequest, }); - await jest.advanceTimersByTimeAsync(20_000); - await enrichPromise; + await enricher.enrichAccount({ accountId, assetsBalance }); - // Batch size is 3; a second batch would be queued without the timeout break. - expect(callCount).toBe(1); - - jest.useRealTimers(); + expect(callSnapRequest).not.toHaveBeenCalled(); + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + }); }); }); }); diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts index 93e8f3fb4e..143a81933b 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -22,13 +22,7 @@ const log = createModuleLogger( export const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; -/** - * Max assets per snap getAccountAssetInfo request. Large batches can terminate - * the Stellar wallet snap on mobile when many trustlines are fetched at once. - */ -const ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE = 1; - -/** Per-batch snap client request timeout (ms). Hung requests must not block apply. */ +/** Snap client request timeout (ms). Hung requests must not block apply. */ const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; // TODO(STELLAR): Replace this chain allowlist with Accounts API-backed enrichment support. @@ -41,6 +35,32 @@ const ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN: Partial< const ENRICHMENT_TIMEOUT = Symbol('enrichmentTimeout'); +type SnapAccountAssetInfoRequest = { + snapId: SnapId; + origin: string; + handler: HandlerType; + request: { + jsonrpc: '2.0'; + method: string; + params: { + accountId: string; + scope: ChainId; + assets: Caip19AssetId[]; + }; + }; +}; + +export type SnapAccountAssetInfoEnricherOptions = { + getSnapIdForChain: (chainId: ChainId) => SnapId | undefined; + callSnapRequest: (request: SnapAccountAssetInfoRequest) => Promise; + log?: typeof log; +}; + +export type EnrichAccountParams = { + accountId: string; + assetsBalance: Record; +}; + function extractChainFromAssetId(assetId: string): ChainId { const parsed = parseCaipAssetType(assetId as CaipAssetType); return parsed.chainId; @@ -77,112 +97,170 @@ export function hasAccountAssetInfoEnrichmentCandidate(params: { return false; } -type SnapAccountAssetInfoRequest = { - snapId: SnapId; - origin: string; - handler: HandlerType; - request: { - jsonrpc: '2.0'; - method: string; - params: { - accountId: string; - scope: ChainId; - assets: Caip19AssetId[]; - }; - }; -}; +/** + * Stateful enricher for Snap `getAccountAssetInfo` responses. + * + * Owns fetch/timeout/apply for a single account. Future in-flight request + * deduplication can key on `${accountId}:${chainId}` inside {@link #fetch}. + */ +export class SnapAccountAssetInfoEnricher { + readonly #getSnapIdForChain: ( + chainId: ChainId, + ) => SnapId | undefined; -async function fetchAccountAssetInfoFromSnap({ - accountId, - snapId, - chainId, - assets, - callSnapRequest, - moduleLog = log, -}: { - accountId: string; - snapId: SnapId; - chainId: ChainId; - assets: Caip19AssetId[]; - callSnapRequest: (request: SnapAccountAssetInfoRequest) => Promise; - moduleLog?: typeof log; -}): Promise { - if (assets.length === 0) { - return undefined; + readonly #callSnapRequest: ( + request: SnapAccountAssetInfoRequest, + ) => Promise; + + readonly #log: typeof log; + + constructor(options: SnapAccountAssetInfoEnricherOptions) { + this.#getSnapIdForChain = options.getSnapIdForChain; + this.#callSnapRequest = options.callSnapRequest; + this.#log = options.log ?? log; } - try { - const request: SnapAccountAssetInfoRequest = { - snapId, - origin: 'metamask', - handler: HandlerType.OnClientRequest, - request: { - jsonrpc: '2.0', - method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, - params: { - accountId, - scope: chainId, - assets, - }, - }, - }; + /** + * Enrich all eligible assets for a single account. + * + * Groups assets by chain, then performs one `getAccountAssetInfo` request + * per (accountId, chainId) and merges the response into `assetsBalance`. + * + * @param params - Account id and that account's balance rows to mutate. + */ + async enrichAccount(params: EnrichAccountParams): Promise { + const { accountId, assetsBalance } = params; + const assetsByChain = this.#getEnrichableAssetsByChain(assetsBalance); - const snapRequest = callSnapRequest(request); - const result = await Promise.race([ - snapRequest, - new Promise((resolve) => - setTimeout(() => resolve(ENRICHMENT_TIMEOUT), ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS), - ), - ]); - - if (result === ENRICHMENT_TIMEOUT) { - snapRequest - .then(() => - moduleLog('Snap account asset info resolved after timeout', { - accountId, - snapId, - chainId, - assetCount: assets.length, - }), - ) - .catch((error) => - moduleLog('Snap account asset info failed after timeout', { - accountId, - snapId, - chainId, - assetCount: assets.length, - error, - }), - ); - return ENRICHMENT_TIMEOUT; + for (const [chainId, assets] of assetsByChain) { + const accountAssetInfo = await this.#fetch(accountId, chainId, assets); + if (!accountAssetInfo) { + continue; + } + this.#apply(accountAssetInfo, assetsBalance); } + } + + /** + * Retrieve enrichment from the Snap for one (accountId, chainId). + * + * Natural place for future in-flight deduplication keyed by + * `${accountId}:${chainId}`. + * + * @param accountId - Account id. + * @param chainId - CAIP-2 chain id. + * @param assets - All enrichable assets on that chain for the account. + * @returns Snap response, or undefined on skip/timeout/error. + */ + async #fetch( + accountId: string, + chainId: ChainId, + assets: Caip19AssetId[], + ): Promise { + if (assets.length === 0) { + return undefined; + } + + const snapId = this.#getSnapIdForChain(chainId); + if (!snapId) { + return undefined; + } + + // Future dedup key: `${accountId}:${chainId}` + + try { + const request: SnapAccountAssetInfoRequest = { + snapId, + origin: 'metamask', + handler: HandlerType.OnClientRequest, + request: { + jsonrpc: '2.0', + method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + params: { + accountId, + scope: chainId, + assets, + }, + }, + }; + + const snapRequest = this.#callSnapRequest(request); + const result = await Promise.race([ + snapRequest, + new Promise((resolve) => + setTimeout( + () => resolve(ENRICHMENT_TIMEOUT), + ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS, + ), + ), + ]); - return result as GetAccountAssetInfoResponse | undefined; - } catch (error) { - moduleLog('Failed to enrich snap account asset info', { - accountId, - snapId, - chainId, - assetCount: assets.length, - error, - }); - return ENRICHMENT_TIMEOUT; + if (result === ENRICHMENT_TIMEOUT) { + snapRequest + .then(() => + this.#log('Snap account asset info resolved after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + }), + ) + .catch((error) => + this.#log('Snap account asset info failed after timeout', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }), + ); + return undefined; + } + + return result as GetAccountAssetInfoResponse | undefined; + } catch (error) { + this.#log('Failed to enrich snap account asset info', { + accountId, + snapId, + chainId, + assetCount: assets.length, + error, + }); + return undefined; + } } -} -export async function enrichAccountAssetInfo(params: { - assetsBalance: NonNullable; - getSnapIdForChain: (chainId: ChainId) => SnapId | undefined; - callSnapRequest: (request: SnapAccountAssetInfoRequest) => Promise; - log?: typeof log; -}): Promise { - const { assetsBalance, getSnapIdForChain, callSnapRequest, log: moduleLog } = - params; + /** + * Merge snap `getAccountAssetInfo` into the account's balance rows. + * + * Pure aside from mutating `assetsBalance`. + * + * @param accountAssetInfo - Enrichment keyed by CAIP-19 asset id. + * @param assetsBalance - Account balance rows to update in place. + */ + #apply( + accountAssetInfo: GetAccountAssetInfoResponse, + assetsBalance: Record, + ): void { + for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { + const row = assetsBalance[assetId as Caip19AssetId]; + if (!row) { + continue; + } - for (const [accountId, accountAssets] of Object.entries(assetsBalance)) { + assetsBalance[assetId as Caip19AssetId] = { + ...row, + accountAssetInfo: assetInfo, + } satisfies AssetBalance; + } + } + + #getEnrichableAssetsByChain( + assetsBalance: Record, + ): Map { const assetsByChain = new Map(); - for (const assetId of Object.keys(accountAssets) as Caip19AssetId[]) { + for (const assetId of Object.keys(assetsBalance) as Caip19AssetId[]) { let chainId: ChainId; try { chainId = extractChainFromAssetId(assetId); @@ -198,49 +276,6 @@ export async function enrichAccountAssetInfo(params: { assetsByChain.set(chainId, assetIds); } - for (const [chainId, assetIds] of assetsByChain) { - const snapId = getSnapIdForChain(chainId); - if (!snapId) { - continue; - } - - for ( - let i = 0; - i < assetIds.length; - i += ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE - ) { - const batch = assetIds.slice( - i, - i + ACCOUNT_ASSET_INFO_SNAP_BATCH_SIZE, - ); - const accountAssetInfo = await fetchAccountAssetInfoFromSnap({ - accountId, - snapId, - chainId, - assets: batch, - callSnapRequest, - moduleLog, - }); - - if ( - accountAssetInfo === ENRICHMENT_TIMEOUT || - accountAssetInfo === undefined - ) { - break; - } - - for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { - const row = accountAssets[assetId as Caip19AssetId]; - if (!row) { - continue; - } - - accountAssets[assetId as Caip19AssetId] = { - ...row, - accountAssetInfo: assetInfo, - } satisfies AssetBalance; - } - } - } + return assetsByChain; } } From 99daf2a36be0121b3524fc1e45b9f98fd500a7ab Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Thu, 9 Jul 2026 16:34:29 +0200 Subject: [PATCH 8/9] chore: dedupe snap request during enrichment --- ...snap-account-asset-info-enrichment.test.ts | 140 ++++++++++++++- .../snap-account-asset-info-enrichment.ts | 162 ++++++++++++++---- 2 files changed, 264 insertions(+), 38 deletions(-) diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts index 59a907a6ee..c716daf17d 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -4,6 +4,7 @@ import { HandlerType } from '@metamask/snaps-utils'; import type { ChainId, Caip19AssetId, DataResponse } from '../types'; import { GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + getAssetsToFetchWithEligibleCustomAssets, hasAccountAssetInfoEnrichmentCandidate, isAccountAssetInfoEnrichmentAvailable, SnapAccountAssetInfoEnricher, @@ -23,16 +24,40 @@ const MOCK_SOL_ASSET = describe('snap-account-asset-info-enrichment', () => { describe('isAccountAssetInfoEnrichmentAvailable', () => { - it('returns true for Stellar chains', () => { + it('returns true for Stellar pubnet', () => { expect(isAccountAssetInfoEnrichmentAvailable(STELLAR_PUBNET)).toBe(true); - expect(isAccountAssetInfoEnrichmentAvailable(STELLAR_TESTNET)).toBe(true); }); it('returns false for non-enrichment chains', () => { + expect(isAccountAssetInfoEnrichmentAvailable(STELLAR_TESTNET)).toBe(false); expect(isAccountAssetInfoEnrichmentAvailable(SOLANA_MAINNET)).toBe(false); }); }); + describe('getAssetsToFetchWithEligibleCustomAssets', () => { + it('adds enrichable custom assets on requested chains', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_STELLAR_USDC_ASSET], + customAssets: [MOCK_STELLAR_ASSET_2], + requestedChainIds: [STELLAR_PUBNET], + }); + + expect(result).toStrictEqual( + expect.arrayContaining([MOCK_STELLAR_USDC_ASSET, MOCK_STELLAR_ASSET_2]), + ); + }); + + it('ignores custom assets on non-enrichment chains', () => { + const result = getAssetsToFetchWithEligibleCustomAssets({ + listedAssets: [MOCK_STELLAR_USDC_ASSET], + customAssets: [MOCK_SOL_ASSET], + requestedChainIds: [STELLAR_PUBNET, SOLANA_MAINNET], + }); + + expect(result).toStrictEqual([MOCK_STELLAR_USDC_ASSET]); + }); + }); + describe('hasAccountAssetInfoEnrichmentCandidate', () => { const accountId = 'mock-account-id'; @@ -184,6 +209,117 @@ describe('snap-account-asset-info-enrichment', () => { }); }); + it('deduplicates concurrent enrichment for the same account and chain', async () => { + const assetsBalance = createAccountAssets(); + let resolveSnapRequest: + | ((value: Record) => void) + | undefined; + const callSnapRequest = jest.fn( + () => + new Promise((resolve) => { + resolveSnapRequest = resolve; + }), + ); + const enricher = createEnricher(callSnapRequest); + + const first = enricher.enrichAccount({ accountId, assetsBalance }); + const second = enricher.enrichAccount({ accountId, assetsBalance }); + + expect(callSnapRequest).toHaveBeenCalledTimes(1); + + resolveSnapRequest?.({ + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + }, + }); + await Promise.all([first, second]); + + expect(callSnapRequest).toHaveBeenCalledTimes(1); + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '1000', + authorized: true, + }, + }); + }); + + it('applies shared fetch result to each distinct assetsBalance object', async () => { + const assetsBalanceA = createAccountAssets(); + const assetsBalanceB = createAccountAssets(); + let resolveSnapRequest: + | ((value: Record) => void) + | undefined; + const callSnapRequest = jest.fn( + () => + new Promise((resolve) => { + resolveSnapRequest = resolve; + }), + ); + const enricher = createEnricher(callSnapRequest); + + const first = enricher.enrichAccount({ + accountId, + assetsBalance: assetsBalanceA, + }); + const second = enricher.enrichAccount({ + accountId, + assetsBalance: assetsBalanceB, + }); + + expect(callSnapRequest).toHaveBeenCalledTimes(1); + + resolveSnapRequest?.({ + [MOCK_STELLAR_USDC_ASSET]: { + limit: '1000', + authorized: true, + }, + }); + await Promise.all([first, second]); + + expect(callSnapRequest).toHaveBeenCalledTimes(1); + expect(assetsBalanceA[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '1000', + authorized: true, + }, + }); + expect(assetsBalanceB[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '1000', + authorized: true, + }, + }); + }); + + it('allows a new enrichment after the previous in-flight operation completes', async () => { + const assetsBalance = createAccountAssets(); + const callSnapRequest = jest + .fn() + .mockResolvedValueOnce({ + [MOCK_STELLAR_USDC_ASSET]: { limit: '1000', authorized: true }, + }) + .mockResolvedValueOnce({ + [MOCK_STELLAR_USDC_ASSET]: { limit: '2000', authorized: false }, + }); + const enricher = createEnricher(callSnapRequest); + + await enricher.enrichAccount({ accountId, assetsBalance }); + await enricher.enrichAccount({ accountId, assetsBalance }); + + expect(callSnapRequest).toHaveBeenCalledTimes(2); + expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ + amount: '25', + accountAssetInfo: { + limit: '2000', + authorized: false, + }, + }); + }); + it('returns balances without accountAssetInfo when enrichment hangs past timeout', async () => { jest.useFakeTimers(); diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts index 143a81933b..b1191adaa8 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -22,7 +22,7 @@ const log = createModuleLogger( export const GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD = 'getAccountAssetInfo'; -/** Snap client request timeout (ms). Hung requests must not block apply. */ +/** Enrichment operation timeout (ms). Hung Snap requests must not block callers. */ const ACCOUNT_ASSET_INFO_SNAP_TIMEOUT_MS = 15_000; // TODO(STELLAR): Replace this chain allowlist with Accounts API-backed enrichment support. @@ -30,7 +30,6 @@ const ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN: Partial< Record > = { 'stellar:pubnet': true, - 'stellar:testnet': true, }; const ENRICHMENT_TIMEOUT = Symbol('enrichmentTimeout'); @@ -70,6 +69,42 @@ export function isAccountAssetInfoEnrichmentAvailable(chainId: ChainId): boolean return ACCOUNT_ASSET_INFO_ENRICHMENT_BY_CHAIN[chainId] === true; } +/** + * Merge listed snap assets with enrichable custom assets on the requested chains. + * + * Custom assets can drop off `listAccountAssets` after trustline deactivation + * (e.g. limit-0 tombstones) while still needing enrichment refresh. + * + * @param params - Listed assets, optional custom assets, and requested chains. + * @returns Deduplicated asset list to fetch balances/enrichment for. + */ +export function getAssetsToFetchWithEligibleCustomAssets(params: { + listedAssets: CaipAssetType[]; + customAssets?: Caip19AssetId[]; + requestedChainIds: ChainId[]; +}): CaipAssetType[] { + const { listedAssets, customAssets, requestedChainIds } = params; + const assetsToFetch = new Set(listedAssets); + + if (customAssets) { + for (const assetId of customAssets) { + try { + const assetChainId = extractChainFromAssetId(assetId); + if ( + requestedChainIds.includes(assetChainId) && + isAccountAssetInfoEnrichmentAvailable(assetChainId) + ) { + assetsToFetch.add(assetId as CaipAssetType); + } + } catch { + // Ignore malformed custom asset ids. + } + } + } + + return [...assetsToFetch]; +} + export function hasAccountAssetInfoEnrichmentCandidate(params: { assetsBalance: NonNullable; getSnapIdForChain: (chainId: ChainId) => SnapId | undefined; @@ -100,8 +135,9 @@ export function hasAccountAssetInfoEnrichmentCandidate(params: { /** * Stateful enricher for Snap `getAccountAssetInfo` responses. * - * Owns fetch/timeout/apply for a single account. Future in-flight request - * deduplication can key on `${accountId}:${chainId}` inside {@link #fetch}. + * Deduplicates the Snap fetch per `${accountId}:${chainId}` so concurrent + * callers share one request, then each caller applies the result to its own + * `assetsBalance` object. */ export class SnapAccountAssetInfoEnricher { readonly #getSnapIdForChain: ( @@ -114,6 +150,11 @@ export class SnapAccountAssetInfoEnricher { readonly #log: typeof log; + readonly #inFlight = new Map< + string, + Promise + >(); + constructor(options: SnapAccountAssetInfoEnricherOptions) { this.#getSnapIdForChain = options.getSnapIdForChain; this.#callSnapRequest = options.callSnapRequest; @@ -123,36 +164,67 @@ export class SnapAccountAssetInfoEnricher { /** * Enrich all eligible assets for a single account. * - * Groups assets by chain, then performs one `getAccountAssetInfo` request - * per (accountId, chainId) and merges the response into `assetsBalance`. + * Groups assets by chain, then runs one deduplicated fetch per + * (accountId, chainId) and applies the result to this caller's balances. * * @param params - Account id and that account's balance rows to mutate. */ async enrichAccount(params: EnrichAccountParams): Promise { const { accountId, assetsBalance } = params; - const assetsByChain = this.#getEnrichableAssetsByChain(assetsBalance); + const assetsByChain = this.#groupAssetsByChain(assetsBalance); for (const [chainId, assets] of assetsByChain) { - const accountAssetInfo = await this.#fetch(accountId, chainId, assets); - if (!accountAssetInfo) { - continue; - } - this.#apply(accountAssetInfo, assetsBalance); + await this.#enrich(accountId, chainId, assets, assetsBalance); } } /** - * Retrieve enrichment from the Snap for one (accountId, chainId). + * Deduplicate the Snap fetch for one (accountId, chainId), then apply to + * this caller's `assetsBalance`. * - * Natural place for future in-flight deduplication keyed by - * `${accountId}:${chainId}`. + * Concurrent callers share the in-flight fetch promise, but each applies + * the response to its own balance object. + * + * @param accountId - Account id. + * @param chainId - CAIP-2 chain id. + * @param assets - All enrichable assets on that chain for the account. + * @param assetsBalance - Account balance rows to update in place. + * @returns Resolves when fetch (shared) and apply (per-caller) complete. + */ + async #enrich( + accountId: string, + chainId: ChainId, + assets: Caip19AssetId[], + assetsBalance: Record, + ): Promise { + const key = `${accountId}:${chainId}`; + let promise = this.#inFlight.get(key); + if (!promise) { + promise = this.#fetchWithTimeout(accountId, chainId, assets).finally( + () => { + this.#inFlight.delete(key); + }, + ); + this.#inFlight.set(key, promise); + } + + const accountAssetInfo = await promise; + if (!accountAssetInfo) { + return; + } + this.#apply(accountAssetInfo, assetsBalance); + } + + /** + * Fetch enrichment from the Snap, racing against a timeout so callers are + * not blocked if the Snap hangs. * * @param accountId - Account id. * @param chainId - CAIP-2 chain id. * @param assets - All enrichable assets on that chain for the account. * @returns Snap response, or undefined on skip/timeout/error. */ - async #fetch( + async #fetchWithTimeout( accountId: string, chainId: ChainId, assets: Caip19AssetId[], @@ -166,25 +238,8 @@ export class SnapAccountAssetInfoEnricher { return undefined; } - // Future dedup key: `${accountId}:${chainId}` - try { - const request: SnapAccountAssetInfoRequest = { - snapId, - origin: 'metamask', - handler: HandlerType.OnClientRequest, - request: { - jsonrpc: '2.0', - method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, - params: { - accountId, - scope: chainId, - assets, - }, - }, - }; - - const snapRequest = this.#callSnapRequest(request); + const snapRequest = this.#fetch(accountId, snapId, chainId, assets); const result = await Promise.race([ snapRequest, new Promise((resolve) => @@ -217,7 +272,7 @@ export class SnapAccountAssetInfoEnricher { return undefined; } - return result as GetAccountAssetInfoResponse | undefined; + return result; } catch (error) { this.#log('Failed to enrich snap account asset info', { accountId, @@ -230,6 +285,41 @@ export class SnapAccountAssetInfoEnricher { } } + /** + * Build and execute a single Snap `getAccountAssetInfo` request. + * + * @param accountId - Account id. + * @param snapId - Snap id for the chain. + * @param chainId - CAIP-2 chain id. + * @param assets - All enrichable assets on that chain for the account. + * @returns Snap response, or undefined. + */ + async #fetch( + accountId: string, + snapId: SnapId, + chainId: ChainId, + assets: Caip19AssetId[], + ): Promise { + const request: SnapAccountAssetInfoRequest = { + snapId, + origin: 'metamask', + handler: HandlerType.OnClientRequest, + request: { + jsonrpc: '2.0', + method: GET_ACCOUNT_ASSET_INFO_CLIENT_METHOD, + params: { + accountId, + scope: chainId, + assets, + }, + }, + }; + + return (await this.#callSnapRequest( + request, + )) as GetAccountAssetInfoResponse | undefined; + } + /** * Merge snap `getAccountAssetInfo` into the account's balance rows. * @@ -255,7 +345,7 @@ export class SnapAccountAssetInfoEnricher { } } - #getEnrichableAssetsByChain( + #groupAssetsByChain( assetsBalance: Record, ): Map { const assetsByChain = new Map(); From 11a4985b4bafc9da5fd45e0d56a6cbc677f29f8b Mon Sep 17 00:00:00 2001 From: Florin Dzeladini Date: Fri, 10 Jul 2026 09:47:33 +0200 Subject: [PATCH 9/9] chore: rename enriched data metadata --- packages/assets-controller/CHANGELOG.md | 4 ++-- .../src/AssetsController.test.ts | 6 +++--- .../src/data-sources/SnapDataSource.test.ts | 2 +- ...snap-account-asset-info-enrichment.test.ts | 20 +++++++++---------- .../snap-account-asset-info-enrichment.ts | 13 ++++++------ packages/assets-controller/src/types.ts | 7 +++++-- 6 files changed, 28 insertions(+), 24 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 0b6ceba901..9bff367784 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,8 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional `snapDataSourceConfig.assetEnrichmentEnabled` getter to `AssetsControllerOptions` so clients can enable or disable snap `getAccountAssetInfo` enrichment at runtime (defaults to disabled) -- Add optional snap account-asset enrichment to fungible balances on `SnapDataSource.fetch` so Stellar trustline fields (`accountAssetInfo`) are included when clients call `getAssets`; balance push events (`AccountsController:accountBalancesUpdated`) pass amounts through without enrichment — clients should call `getAssets` after trustline-related snap events (e.g. `AccountsController:accountAssetListUpdated`) +- Add optional `snapDataSourceConfig.assetEnrichmentEnabled` getter to `AssetsControllerOptions` so clients can enable or disable snap `getAccountAssetInfo` enrichment at runtime (defaults to disabled) ([#9392](https://github.com/MetaMask/core/pull/9217)) +- Add optional snap account-asset enrichment to balances in `SnapDataSource.fetch`, so `getAssets` includes Stellar trustline metadata. Balance updates skip enrichment; use `getAssets` after trustline events. ([#9392](https://github.com/MetaMask/core/pull/9217)) ### Changed diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 9442987fd9..e8ff20615b 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -1687,13 +1687,13 @@ describe('AssetsController', () => { }); }); - it('preserves accountAssetInfo enrichment when merge update changes amount', async () => { + it('preserves metadata enrichment when merge update changes amount', async () => { const initialState: Partial = { assetsBalance: { [MOCK_ACCOUNT_ID]: { [MOCK_ASSET_ID]: { amount: '1', - accountAssetInfo: { limit: '1000', authorized: true }, + metadata: { limit: '1000', authorized: true }, }, }, }, @@ -1716,7 +1716,7 @@ describe('AssetsController', () => { controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID], ).toStrictEqual({ amount: '2', - accountAssetInfo: { limit: '1000', authorized: true }, + metadata: { limit: '1000', authorized: true }, }); }); }); diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index ef9710afe9..03628e9d3c 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -602,7 +602,7 @@ describe('SnapDataSource', () => { response.assetsBalance?.['mock-account-id']?.[MOCK_STELLAR_USDC_ASSET], ).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '1000', authorized: true, sponsored: false, diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts index c716daf17d..c4c31a76b3 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.test.ts @@ -134,7 +134,7 @@ describe('snap-account-asset-info-enrichment', () => { }); } - it('enriches Stellar assets with accountAssetInfo', async () => { + it('enriches Stellar assets with metadata', async () => { const assetsBalance = createAccountAssets(); const callSnapRequest = jest.fn().mockResolvedValue({ [MOCK_STELLAR_USDC_ASSET]: { @@ -163,7 +163,7 @@ describe('snap-account-asset-info-enrichment', () => { }); expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '1000', authorized: true, sponsored: false, @@ -201,11 +201,11 @@ describe('snap-account-asset-info-enrichment', () => { ); expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '10', - accountAssetInfo: { limit: '1000', authorized: true }, + metadata: { limit: '1000', authorized: true }, }); expect(assetsBalance[MOCK_STELLAR_ASSET_2]).toStrictEqual({ amount: '20', - accountAssetInfo: { limit: '500', authorized: false }, + metadata: { limit: '500', authorized: false }, }); }); @@ -238,7 +238,7 @@ describe('snap-account-asset-info-enrichment', () => { expect(callSnapRequest).toHaveBeenCalledTimes(1); expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '1000', authorized: true, }, @@ -281,14 +281,14 @@ describe('snap-account-asset-info-enrichment', () => { expect(callSnapRequest).toHaveBeenCalledTimes(1); expect(assetsBalanceA[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '1000', authorized: true, }, }); expect(assetsBalanceB[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '1000', authorized: true, }, @@ -313,14 +313,14 @@ describe('snap-account-asset-info-enrichment', () => { expect(callSnapRequest).toHaveBeenCalledTimes(2); expect(assetsBalance[MOCK_STELLAR_USDC_ASSET]).toStrictEqual({ amount: '25', - accountAssetInfo: { + metadata: { limit: '2000', authorized: false, }, }); }); - it('returns balances without accountAssetInfo when enrichment hangs past timeout', async () => { + it('returns balances without metadata when enrichment hangs past timeout', async () => { jest.useFakeTimers(); const assetsBalance = createAccountAssets(); @@ -341,7 +341,7 @@ describe('snap-account-asset-info-enrichment', () => { expect( ( assetsBalance[MOCK_STELLAR_USDC_ASSET] as Record - )?.accountAssetInfo, + )?.metadata, ).toBeUndefined(); jest.useRealTimers(); diff --git a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts index b1191adaa8..ee72191302 100644 --- a/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts +++ b/packages/assets-controller/src/data-sources/snap-account-asset-info-enrichment.ts @@ -1,4 +1,4 @@ -// TODO(STELLAR): This helper is a temporary bridge for Snap-provided accountAssetInfo. +// TODO(STELLAR): This helper is a temporary bridge for Snap-provided balance metadata. // Remove it once the Accounts API supports account-asset enrichment directly. import type { CaipAssetType } from '@metamask/keyring-api'; @@ -321,18 +321,19 @@ export class SnapAccountAssetInfoEnricher { } /** - * Merge snap `getAccountAssetInfo` into the account's balance rows. + * Merge snap `getAccountAssetInfo` into the account's balance rows as + * `metadata`. * * Pure aside from mutating `assetsBalance`. * - * @param accountAssetInfo - Enrichment keyed by CAIP-19 asset id. + * @param enrichment - Enrichment keyed by CAIP-19 asset id. * @param assetsBalance - Account balance rows to update in place. */ #apply( - accountAssetInfo: GetAccountAssetInfoResponse, + enrichment: GetAccountAssetInfoResponse, assetsBalance: Record, ): void { - for (const [assetId, assetInfo] of Object.entries(accountAssetInfo)) { + for (const [assetId, assetInfo] of Object.entries(enrichment)) { const row = assetsBalance[assetId as Caip19AssetId]; if (!row) { continue; @@ -340,7 +341,7 @@ export class SnapAccountAssetInfoEnricher { assetsBalance[assetId as Caip19AssetId] = { ...row, - accountAssetInfo: assetInfo, + metadata: assetInfo, } satisfies AssetBalance; } } diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index 0fc27c1a8d..65bd41b160 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -292,8 +292,11 @@ export type GetAccountAssetInfoResponse = Record< export type FungibleAssetBalance = { /** Raw balance amount as string (e.g., "1000000000" for 1000 USDC) */ amount: string; - /** Optional per-asset info from snap account-asset enrichment (e.g. Stellar trustlines). */ - accountAssetInfo?: AccountAssetInfo; + /** + * Optional per-asset account metadata from snap enrichment + * (e.g. Stellar trustline fields). Distinct from `assetsInfo` token metadata. + */ + metadata?: AccountAssetInfo; }; /**