fix(assets-controller): improve SnapDataSource chain discovery and balance delivery

eslint-suppressions.json

@@ -183,11 +183,6 @@
       "count": 2
     }
   },
-  "packages/assets-controller/src/data-sources/SnapDataSource.ts": {
-    "no-restricted-syntax": {
-      "count": 2
-    }
-  },
   "packages/assets-controller/src/data-sources/StakedBalanceDataSource.ts": {
     "no-restricted-syntax": {
       "count": 2

packages/assets-controller/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- SnapDataSource now subscribes to `PermissionController:stateChange` instead of deprecated `PermissionController:stateChange`. Hosts that restrict or delegate events into the `AssetsController` messenger must delegate `PermissionController:stateChange`; delegating only `stateChange` will prevent permission-driven snap rediscovery. ([#8857](https://github.com/MetaMask/core/pull/8857))
 - Bump `@metamask/transaction-controller` from `^65.3.0` to `^66.0.0` ([#8796](https://github.com/MetaMask/core/pull/8796), [#8848](https://github.com/MetaMask/core/pull/8848))
 - Bump `@metamask/core-backend` from `^6.2.2` to `^6.3.0` ([#8813](https://github.com/MetaMask/core/pull/8813))
 - Bump `@metamask/phishing-controller` from `^17.1.2` to `^17.2.0` ([#8819](https://github.com/MetaMask/core/pull/8819))
@@ -17,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **SnapDataSource:** Read supported CAIP-2 chain IDs from `endowment:keyring` when `endowment:assets` has no `chainIds` caveat (common for some wallet snaps); retry discovery when a balance push arrives before `activeChains` is populated; accept balance payloads while discovery is still empty so updates are not dropped; register subscriptions in that window so `AccountsController:accountBalancesUpdated` can reach `onAssetsUpdate`. ([#8857](https://github.com/MetaMask/core/pull/8857))
 - Non-EVM assets with a `slip44` asset namespace (e.g. Bitcoin, Solana native, TRON) are now correctly typed as `native` instead of `erc20` in `assetsInfo` ([#8811](https://github.com/MetaMask/core/pull/8811))
 - Solana SPL tokens (CAIP-19 `solana:.../token:<address>`) are now correctly typed as `spl` instead of `erc20` in `assetsInfo` ([#8811](https://github.com/MetaMask/core/pull/8811))
 

packages/assets-controller/src/AssetsController.ts

@@ -37,10 +37,7 @@ import type {
   NetworkEnablementControllerEvents,
   NetworkEnablementControllerState,
 } from '@metamask/network-enablement-controller';
-import type {
-  GetPermissions,
-  PermissionControllerStateChange,
-} from '@metamask/permission-controller';
+import type { GetPermissions } from '@metamask/permission-controller';
 import { PhishingControllerBulkScanTokensAction } from '@metamask/phishing-controller';
 import type { PreferencesControllerStateChangeEvent } from '@metamask/preferences-controller';
 import type {
@@ -79,7 +76,7 @@ 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 { SnapDataSourceAllowedEvents } from './data-sources/SnapDataSource';
 import { SnapDataSource } from './data-sources/SnapDataSource';
 import type { StakedBalanceDataSourceConfig } from './data-sources/StakedBalanceDataSource';
 import { StakedBalanceDataSource } from './data-sources/StakedBalanceDataSource';
@@ -335,8 +332,7 @@ type AllowedEvents =
   // StakedBalanceDataSource
   | NetworkEnablementControllerEvents
   // SnapDataSource
-  | AccountsControllerAccountBalancesUpdatedEvent
-  | PermissionControllerStateChange
+  | SnapDataSourceAllowedEvents
   // BackendWebsocketDataSource
   | BackendWebSocketServiceEvents;
 
@@ -2565,9 +2561,13 @@ export class AssetsController extends BaseController<
     accounts: InternalAccount[],
     chainIds: ChainId[],
   ): void {
+    // Snap data source handles non-EVM chains that are never in #enabledChains.
+    // Include its active chains so snap-backed accounts (bip122, solana, tron…)
+    // appear in chainToAccounts and receive a subscription.
+    const snapActiveChains = this.#snapDataSource.getActiveChainsSync();
     const chainToAccounts = this.#buildChainToAccountsMap(
       accounts,
-      new Set(chainIds),
+      new Set([...chainIds, ...snapActiveChains]),
     );
     const remainingChains = new Set(chainToAccounts.keys());
     // When basic functionality is on, use all balance data sources; when off, RPC only.

packages/assets-controller/src/data-sources/SnapDataSource.test.ts

@@ -146,6 +146,31 @@ function createMockPermissions(
   } as unknown as SubjectPermissions<PermissionConstraint>;
 }
 
+/**
+ * Mock permissions where chainIds exist only on `endowment:keyring` (no assets permission).
+ *
+ * @param chainIds - Chain IDs for the caveat.
+ * @returns Permissions object.
+ */
+function createMockPermissionsKeyringOnly(
+  chainIds: ChainId[],
+): SubjectPermissions<PermissionConstraint> {
+  return {
+    [KEYRING_PERMISSION]: {
+      id: 'mock-keyring-permission-id',
+      parentCapability: KEYRING_PERMISSION,
+      invoker: 'test',
+      date: Date.now(),
+      caveats: [
+        {
+          type: 'chainIds',
+          value: chainIds,
+        },
+      ],
+    },
+  } as unknown as SubjectPermissions<PermissionConstraint>;
+}
+
 /**
  * Creates a mock handler for SnapController:handleRequest
  *
@@ -171,13 +196,26 @@ function createMockHandleRequest(
 
 function setupController(
   options: {
-    installedSnaps?: Record<string, { version: string; chainIds?: ChainId[] }>;
+    installedSnaps?: Record<
+      string,
+      {
+        version: string;
+        chainIds?: ChainId[];
+        keyringOnlyChainIds?: ChainId[];
+      }
+    >;
     accountAssets?: string[];
     balances?: Record<string, { amount: string; unit: string }>;
     configuredNetworks?: ChainId[];
+    getRunnableSnapsImplementation?: () => unknown;
   } = {},
 ): SetupResult {
-  const { installedSnaps = {}, accountAssets = [], balances = {} } = options;
+  const {
+    installedSnaps = {},
+    accountAssets = [],
+    balances = {},
+    getRunnableSnapsImplementation,
+  } = options;
 
   const rootMessenger = new Messenger<MockAnyNamespace, AllActions, AllEvents>({
     namespace: MOCK_ANY_NAMESPACE,
@@ -200,7 +238,10 @@ function setupController(
       'SnapController:handleRequest',
       'PermissionController:getPermissions',
     ],
-    events: ['AccountsController:accountBalancesUpdated'],
+    events: [
+      'AccountsController:accountBalancesUpdated',
+      'PermissionController:stateChange',
+    ],
   });
 
   const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined);
@@ -218,9 +259,12 @@ function setupController(
   );
 
   // Register SnapController action handlers
-  const mockGetRunnableSnaps = jest
-    .fn()
-    .mockReturnValue(snapsForGetRunnableSnaps);
+  const mockGetRunnableSnaps = jest.fn(
+    getRunnableSnapsImplementation ??
+      ((): typeof snapsForGetRunnableSnaps => {
+        return snapsForGetRunnableSnaps;
+      }),
+  );
   rootMessenger.registerActionHandler(
     'SnapController:getRunnableSnaps',
     mockGetRunnableSnaps,
@@ -236,9 +280,12 @@ function setupController(
   // Returns permissions with chainIds caveat based on installed snaps config
   const mockGetPermissions = jest.fn().mockImplementation((snapId: string) => {
     const snapConfig = installedSnaps[snapId];
-    if (snapConfig?.chainIds) {
+    if (snapConfig?.chainIds?.length) {
       return createMockPermissions(snapConfig.chainIds);
     }
+    if (snapConfig?.keyringOnlyChainIds?.length) {
+      return createMockPermissionsKeyringOnly(snapConfig.keyringOnlyChainIds);
+    }
     return undefined;
   });
   rootMessenger.registerActionHandler(
@@ -409,6 +456,166 @@ describe('SnapDataSource', () => {
     cleanup();
   });
 
+  it('discovers chain IDs from endowment:keyring when assets permission has no chainIds caveat', async () => {
+    const { controller, cleanup } = setupController({
+      installedSnaps: {
+        [SOLANA_SNAP_ID]: {
+          version: '1.0.0',
+          keyringOnlyChainIds: [SOLANA_MAINNET, TRON_MAINNET],
+        },
+      },
+    });
+    await new Promise(process.nextTick);
+
+    const chains = await controller.getActiveChains();
+    expect(chains).toStrictEqual(
+      expect.arrayContaining([SOLANA_MAINNET, TRON_MAINNET]),
+    );
+
+    cleanup();
+  });
+
+  it('forwards accountBalancesUpdated while snap discovery has not completed (fail-open)', async () => {
+    const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined);
+    const { controller, triggerBalancesUpdated, cleanup } = setupController({
+      getRunnableSnapsImplementation: () => {
+        throw new Error('SnapController not ready');
+      },
+    });
+    await new Promise(process.nextTick);
+
+    await controller.subscribe({
+      request: createDataRequest({
+        chainIds: [SOLANA_MAINNET],
+      }),
+      subscriptionId: 'sub-fail-open',
+      isUpdate: false,
+      onAssetsUpdate: assetsUpdateHandler,
+    });
+
+    triggerBalancesUpdated({
+      balances: {
+        'mock-account-id': {
+          [MOCK_SOL_ASSET]: { amount: '1', unit: 'SOL' },
+        },
+      },
+    });
+
+    expect(assetsUpdateHandler).toHaveBeenCalledTimes(1);
+    expect(assetsUpdateHandler).toHaveBeenCalledWith(
+      expect.objectContaining({
+        updateMode: 'merge',
+        assetsBalance: {
+          'mock-account-id': {
+            [MOCK_SOL_ASSET]: { amount: '1' },
+          },
+        },
+      }),
+    );
+
+    cleanup();
+  });
+
+  it('does not fail-open after discovery completes with no supported chains', async () => {
+    const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined);
+    const { controller, triggerBalancesUpdated, cleanup } = setupController({
+      installedSnaps: {
+        [SOLANA_SNAP_ID]: { version: '1.0.0' },
+      },
+    });
+    await new Promise(process.nextTick);
+
+    await controller.subscribe({
+      request: createDataRequest({
+        chainIds: [SOLANA_MAINNET],
+      }),
+      subscriptionId: 'sub-no-discovered-chains',
+      isUpdate: false,
+      onAssetsUpdate: assetsUpdateHandler,
+    });
+
+    triggerBalancesUpdated({
+      balances: {
+        'mock-account-id': {
+          [MOCK_SOL_ASSET]: { amount: '1', unit: 'SOL' },
+        },
+      },
+    });
+
+    await new Promise(process.nextTick);
+
+    expect(assetsUpdateHandler).not.toHaveBeenCalled();
+
+    cleanup();
+  });
+
+  it('keeps last discovered chains when rediscovery fails', async () => {
+    const assetsUpdateHandler = jest.fn().mockResolvedValue(undefined);
+    let shouldFailRediscovery = false;
+    const { controller, messenger, triggerBalancesUpdated, cleanup } =
+      setupController({
+        installedSnaps: {
+          [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] },
+        },
+        getRunnableSnapsImplementation: () => {
+          if (shouldFailRediscovery) {
+            throw new Error('SnapController not ready');
+          }
+          return [
+            {
+              id: SOLANA_SNAP_ID,
+              version: '1.0.0',
+              enabled: true,
+              blocked: false,
+            },
+          ];
+        },
+      });
+    await new Promise(process.nextTick);
+
+    expect(await controller.getActiveChains()).toStrictEqual([SOLANA_MAINNET]);
+
+    shouldFailRediscovery = true;
+    messenger.publish('PermissionController:stateChange', { subjects: {} }, []);
+
+    expect(await controller.getActiveChains()).toStrictEqual([SOLANA_MAINNET]);
+
+    await controller.subscribe({
+      request: createDataRequest({
+        chainIds: [SOLANA_MAINNET],
+      }),
+      subscriptionId: 'sub-after-failed-rediscovery',
+      isUpdate: false,
+      onAssetsUpdate: assetsUpdateHandler,
+    });
+
+    assetsUpdateHandler.mockClear();
+    triggerBalancesUpdated({
+      balances: {
+        'mock-account-id': {
+          [MOCK_SOL_ASSET]: { amount: '1', unit: 'SOL' },
+          ['eip155:1/slip44:60' as Caip19AssetId]: {
+            amount: '1',
+            unit: 'ETH',
+          },
+        },
+      },
+    });
+
+    expect(assetsUpdateHandler).toHaveBeenCalledTimes(1);
+    expect(assetsUpdateHandler).toHaveBeenCalledWith(
+      expect.objectContaining({
+        assetsBalance: {
+          'mock-account-id': {
+            [MOCK_SOL_ASSET]: { amount: '1' },
+          },
+        },
+      }),
+    );
+
+    cleanup();
+  });
+
   it('fetch returns empty response for accounts without snap metadata', async () => {
     const { controller, cleanup } = setupController({
       installedSnaps: {
@@ -922,7 +1129,10 @@ describe('SnapDataSource', () => {
         'SnapController:handleRequest',
         'PermissionController:getPermissions',
       ],
-      events: ['AccountsController:accountBalancesUpdated'],
+      events: [
+        'AccountsController:accountBalancesUpdated',
+        'PermissionController:stateChange',
+      ],
     });
     rootMessenger.registerActionHandler(
       'SnapController:getRunnableSnaps',