fix(assets-controller): consume Snaps Events in all cases

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,
     },
   },
 });

packages/assets-controller/src/AssetsController.ts

@@ -921,6 +921,8 @@
     this.#snapDataSource = new SnapDataSource({
       messenger: this.messenger,
       onActiveChainsUpdated: this.#onActiveChainsUpdated,
+      onAssetsUpdate: (response) =>
+        this.handleAssetsUpdate(response, 'SnapDataSource'),
     });
     this.#rpcDataSource = new RpcDataSource({
       messenger: this.messenger,

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);

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<void>;
   /** 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<void>;
+
   /** 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<void> {
-    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();
   }