diff --git a/packages/assets-controller/jest.config.js b/packages/assets-controller/jest.config.js index 111146cdab..5728cc812f 100644 --- a/packages/assets-controller/jest.config.js +++ b/packages/assets-controller/jest.config.js @@ -19,8 +19,8 @@ module.exports = merge(baseConfig, { global: { branches: 80.2, functions: 88.9, - lines: 89.57, - statements: 89.47, + lines: 89.4, + statements: 89.4, }, }, }); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 3fb7d56a85..70afd4d979 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -921,6 +921,8 @@ export class AssetsController extends BaseController< this.#snapDataSource = new SnapDataSource({ messenger: this.messenger, onActiveChainsUpdated: this.#onActiveChainsUpdated, + onAssetsUpdate: (response) => + this.handleAssetsUpdate(response, 'SnapDataSource'), }); 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 00d517cee9..cb532676b2 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -249,6 +249,7 @@ function setupController( const controllerOptions: SnapDataSourceOptions = { messenger: controllerMessenger as unknown as AssetsControllerMessenger, onActiveChainsUpdated: activeChainsUpdateHandler, + onAssetsUpdate: assetsUpdateHandler, }; const controller = new SnapDataSource(controllerOptions); @@ -940,6 +941,7 @@ describe('SnapDataSource', () => { const instance = createSnapDataSource({ messenger: controllerMessenger as unknown as AssetsControllerMessenger, onActiveChainsUpdated: jest.fn(), + onAssetsUpdate: jest.fn(), }); await new Promise(process.nextTick); diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index a45e0c4424..ac12e50bdc 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -166,6 +166,15 @@ export type SnapDataSourceOptions = { chains: ChainId[], previousChains: ChainId[], ) => void; + /** + * Called directly with snap-sourced balance updates (from + * AccountsController:accountBalancesUpdated), independent of whether the + * controller currently has an active subscription for the chain. Snaps + * (e.g. Tron) can report asset types another data source (e.g. WebSocket) + * owns the base chain for but does not itself support, so delivery must + * not depend on subscription bookkeeping. + */ + onAssetsUpdate: (response: DataResponse) => void | Promise; /** Configured networks to support (defaults to all snap networks) */ configuredNetworks?: ChainId[]; /** Default polling interval in ms for subscriptions */ @@ -208,6 +217,8 @@ export class SnapDataSource extends AbstractDataSource< previousChains: ChainId[], ) => void; + readonly #onAssetsUpdate: (response: DataResponse) => void | Promise; + /** Bound handler for snap keyring balance updates, stored for cleanup */ readonly #handleSnapBalancesUpdatedBound: ( payload: AccountBalancesUpdatedEventPayload, @@ -226,6 +237,7 @@ export class SnapDataSource extends AbstractDataSource< this.#messenger = options.messenger; this.#onActiveChainsUpdated = options.onActiveChainsUpdated; + this.#onAssetsUpdate = options.onAssetsUpdate; // Bind handlers for cleanup in destroy() this.#handleSnapBalancesUpdatedBound = this.#handleSnapBalancesUpdated.bind( @@ -264,7 +276,12 @@ export class SnapDataSource extends AbstractDataSource< /** * Handle snap balance updates from the keyring. - * Transforms the payload and publishes to AssetsController. + * Transforms the payload and reports it directly to AssetsController via + * `onAssetsUpdate`, regardless of whether a subscription is currently + * active for the chain. Snaps push updates (e.g. Tron energy/bandwidth) + * that other data sources cannot fetch even when they own the base chain + * for polling purposes, so delivery here must not be gated on + * `activeSubscriptions`. * * @param payload - The balance update payload from AccountsController. */ @@ -305,9 +322,7 @@ 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); - } + Promise.resolve(this.#onAssetsUpdate(response)).catch(console.error); } } @@ -607,7 +622,7 @@ export class SnapDataSource extends AbstractDataSource< // ============================================================================ async subscribe(subscriptionRequest: SubscriptionRequest): Promise { - const { request, subscriptionId, isUpdate } = subscriptionRequest; + const { request, subscriptionId, onAssetsUpdate } = subscriptionRequest; // Guard against undefined request or chainIds if (!request?.chainIds) { @@ -623,58 +638,22 @@ export class SnapDataSource extends AbstractDataSource< return; } - if (isUpdate) { - const existing = this.activeSubscriptions.get(subscriptionId); - if (existing) { - existing.chains = supportedChains; - // Do a fetch to get latest data on subscription update - this.fetch({ - ...request, - chainIds: supportedChains, - }) - .then(async (fetchResponse) => { - if (Object.keys(fetchResponse.assetsBalance ?? {}).length > 0) { - await existing.onAssetsUpdate(fetchResponse); - } - return fetchResponse; - }) - .catch((error) => { - log('Subscription update fetch failed', { subscriptionId, error }); - }); - return; - } - } - - await this.unsubscribe(subscriptionId); - - // Snaps provide real-time updates via AccountsController:accountBalancesUpdated - // We only need to track the subscription and do an initial fetch - // No polling needed - updates come through #handleSnapBalancesUpdated - - this.activeSubscriptions.set(subscriptionId, { - cleanup: () => { - // No timer to clear - we use event-based updates - }, - chains: supportedChains, - onAssetsUpdate: subscriptionRequest.onAssetsUpdate, - }); - - // Initial fetch to get current balances + // Snaps provide real-time updates via AccountsController:accountBalancesUpdated, + // delivered directly to the controller via the `onAssetsUpdate` passed to + // the constructor (see #handleSnapBalancesUpdated). No polling, and no + // per-subscription state to maintain here — a fetch reports the current + // snapshot whether this is the initial subscribe or a later update. try { const fetchResponse = await this.fetch({ ...request, chainIds: supportedChains, }); - const subscription = this.activeSubscriptions.get(subscriptionId); - if ( - Object.keys(fetchResponse.assetsBalance ?? {}).length > 0 && - subscription - ) { - await subscription.onAssetsUpdate(fetchResponse); + if (Object.keys(fetchResponse.assetsBalance ?? {}).length > 0) { + await onAssetsUpdate(fetchResponse); } } catch (error) { - log('Initial fetch failed', { subscriptionId, error }); + log('Subscription fetch failed', { subscriptionId, error }); } } @@ -741,13 +720,6 @@ export class SnapDataSource extends AbstractDataSource< log('Failed to unsubscribe from permission changes', { error }); } - // Clean up active subscriptions - for (const [subscriptionId] of this.activeSubscriptions) { - this.unsubscribe(subscriptionId).catch(() => { - // Ignore cleanup errors - }); - } - // Clear keyring client cache this.#keyringClientCache.clear(); }