feat(WPN-1476): route Solana asset reads via migration stages (Lane A - PR 4a)

.yarnrc.yml

@@ -17,8 +17,10 @@ nodeLinker: node-modules
 npmMinimalAgeGate: 4320
 
 npmPreapprovedPackages:
+  - "@metamask/assets-controller"
   - "@metamask/keyring-api"
   - "@metamask/keyring-snap-sdk"
+  - "@metamask-previews/snap-networks-utils"
 
 plugins:
   - path: .yarn/plugins/@yarnpkg/plugin-allow-scripts.cjs

package.json

@@ -63,7 +63,7 @@
   },
   "resolutions": {
     "@metamask/snaps-execution-environments": "11.2.0",
-    "@metamask/snaps-sdk": "^11.1.1",
+    "@metamask/snaps-sdk": "^11.2.0",
     "@types/react": "18.2.4",
     "@types/react-dom": "18.2.4"
   },

packages/snap/jest.setup.ts

@@ -1,14 +1,46 @@
 import { jest } from '@jest/globals';
+import type { SimulationUserOptions } from '@metamask/snaps-simulation';
 import BigNumber from 'bignumber.js';
 import dotenv from 'dotenv';
 
+import { registerCoreAssetsControllerHandlers } from './src/core/test/helpers/registerCoreAssetsControllerHandlers';
 import logger from './src/core/utils/logger';
 
 dotenv.config();
 
 // Lowest precision we ever go for: MicroLamports represented in Sol amount
 BigNumber.config({ EXPONENTIAL_AT: 16 });
 
+type SnapsTestEnvironment = {
+  installSnap: (
+    snapId?: string,
+    options?: { options?: SimulationUserOptions },
+  ) => Promise<{
+    controllerMessenger: Parameters<
+      typeof registerCoreAssetsControllerHandlers
+    >[0];
+  }>;
+};
+
+const { snapsEnvironment } = globalThis as {
+  snapsEnvironment?: SnapsTestEnvironment;
+};
+
+if (snapsEnvironment) {
+  const originalInstallSnap =
+    snapsEnvironment.installSnap.bind(snapsEnvironment);
+  jest
+    .spyOn(snapsEnvironment, 'installSnap')
+    .mockImplementation(async (snapId, options = {}) => {
+      const installed = await originalInstallSnap(snapId, options);
+      registerCoreAssetsControllerHandlers(
+        installed.controllerMessenger,
+        options.options ?? {},
+      );
+      return installed;
+    });
+}
+
 // Mock the console methods
 jest.spyOn(logger, 'log').mockImplementation(() => {
   /* no-op */

packages/snap/package.json

@@ -39,22 +39,28 @@
     "prepublishOnly": "mm-snap manifest",
     "serve": "mm-snap serve",
     "start": "node scripts/update-manifest-local.js && concurrently \"mm-snap watch\" \"yarn locale:build:watch\"",
-    "test": "jest --passWithNoTests --runInBand",
+    "test": "ENVIRONMENT=test yarn build && jest --passWithNoTests --runInBand",
     "test:core": "jest src/core --passWithNoTests",
     "test:core:watch": "yarn test:core --watch",
     "test:features": "jest src/features --passWithNoTests --runInBand",
     "test:features:watch": "yarn test:features --watch",
     "test:watch": "yarn test --watch"
   },
+  "dependencies": {
+    "@metamask-previews/snap-networks-utils": "0.0.0-preview-8589132"
+  },
   "devDependencies": {
     "@jest/globals": "^29.5.0",
+    "@metamask/assets-controller": "13.0.0",
     "@metamask/auto-changelog": "4.0.0",
     "@metamask/key-tree": "9.1.2",
     "@metamask/keyring-api": "^23.5.0",
     "@metamask/keyring-snap-sdk": "^9.2.0",
+    "@metamask/messenger": "^2.0.0",
+    "@metamask/remote-feature-flag-controller": "5.0.0",
     "@metamask/snaps-cli": "^8.4.1",
     "@metamask/snaps-jest": "^10.2.0",
-    "@metamask/snaps-sdk": "^11.1.1",
+    "@metamask/snaps-sdk": "^11.2.0",
     "@metamask/superstruct": "^3.1.0",
     "@metamask/utils": "^11.9.0",
     "@noble/ed25519": "2.1.0",

packages/snap/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snap-solana-wallet.git"
   },
   "source": {
-    "shasum": "IFUCW/5nFjm2PtAdzVtjtdmN9x8sVAYMvTCnp1fivHU=",
+    "shasum": "KRhMEDiB6VpAZciUc7txZxIbo2D1q3UMDo/9givv7X8=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",
@@ -88,8 +88,16 @@
     "snap_manageAccounts": {},
     "snap_manageState": {},
     "snap_dialog": {},
-    "snap_getPreferences": {}
+    "snap_getPreferences": {},
+    "endowment:messenger": {
+      "actions": [
+        "AssetsController:getAccountAssetByID",
+        "AssetsController:getAccountAssetsByIDs",
+        "AssetsController:getAccountAssetsByScope",
+        "RemoteFeatureFlagController:getState"
+      ]
+    }
   },
-  "platformVersion": "11.1.1",
+  "platformVersion": "11.2.0",
   "manifestVersion": "0.1"
 }

packages/snap/src/core/handlers/onCronjob/backgroundEvents/refreshSend.test.tsx

@@ -1,4 +1,10 @@
-import { assetsService, priceApiClient, state } from '../../../../snapContext';
+import {
+  accountsService,
+  assetsService,
+  configProvider,
+  priceApiClient,
+  state,
+} from '../../../../snapContext';
 import { KnownCaip19Id } from '../../../constants/solana';
 import { trackError } from '../../../utils/errors';
 import {
@@ -33,9 +39,15 @@ jest.mock('../../../../features/send/Send', () => ({
 }));
 
 jest.mock('../../../../snapContext', () => ({
-  assetsService: {
+  accountsService: {
     getAll: jest.fn(),
   },
+  assetsService: {
+    getAccountAssetsByScope: jest.fn(),
+  },
+  configProvider: {
+    getActiveNetworks: jest.fn(),
+  },
   priceApiClient: {
     getMultipleSpotPrices: jest.fn(),
   },
@@ -50,7 +62,13 @@ const setupTest = () => {
     request: jest.fn(),
   };
 
-  (assetsService.getAll as jest.Mock).mockResolvedValue([
+  (accountsService.getAll as jest.Mock).mockResolvedValue([
+    { id: 'account-1' },
+  ]);
+  (configProvider.getActiveNetworks as jest.Mock).mockResolvedValue([
+    'solana:mainnet',
+  ]);
+  (assetsService.getAccountAssetsByScope as jest.Mock).mockResolvedValue([
     {
       assetType: KnownCaip19Id.SolMainnet,
     },

packages/snap/src/core/handlers/onCronjob/backgroundEvents/refreshSend.tsx

@@ -3,7 +3,13 @@ import type { OnCronjobHandler } from '@metamask/snaps-sdk';
 import { DEFAULT_SEND_CONTEXT } from '../../../../features/send/render';
 import { Send } from '../../../../features/send/Send';
 import type { SendContext } from '../../../../features/send/types';
-import { assetsService, priceApiClient, state } from '../../../../snapContext';
+import {
+  assetsService,
+  configProvider,
+  priceApiClient,
+  state,
+  accountsService,
+} from '../../../../snapContext';
 import type { UnencryptedStateValue } from '../../../services/state/State';
 import { trackError } from '../../../utils/errors';
 import {
@@ -19,13 +25,25 @@ export const refreshSend: OnCronjobHandler = async () => {
 
   logger.info(`Background event triggered`);
 
-  const [assets, mapInterfaceNameToId, preferences] = await Promise.all([
-    assetsService.getAll(),
-    state.getKey<UnencryptedStateValue['mapInterfaceNameToId']>(
-      'mapInterfaceNameToId',
-    ),
-    getPreferences().catch(() => DEFAULT_SEND_CONTEXT.preferences),
-  ]);
+  const [accounts, activeNetworks, mapInterfaceNameToId, preferences] =
+    await Promise.all([
+      accountsService.getAll(),
+      configProvider.getActiveNetworks(),
+      state.getKey<UnencryptedStateValue['mapInterfaceNameToId']>(
+        'mapInterfaceNameToId',
+      ),
+      getPreferences().catch(() => DEFAULT_SEND_CONTEXT.preferences),
+    ]);
+
+  const assets = (
+    await Promise.all(
+      accounts.flatMap((account) =>
+        activeNetworks.map((network) =>
+          assetsService.getAccountAssetsByScope(network, account.id),
+        ),
+      ),
+    )
+  ).flat();
 
   const assetTypes = assets.flatMap((asset) => asset.assetType);
 

packages/snap/src/core/handlers/onKeyringRequest/Keyring.test.ts

@@ -103,7 +103,8 @@ describe('SolanaKeyring', () => {
     mockAssetsService = {
       fetch: jest.fn().mockResolvedValue(MOCK_ASSET_ENTITIES),
       saveMany: jest.fn(),
-      findByAccount: jest.fn(),
+      getAccountAssetsForAllActiveScopes: jest.fn(),
+      getAccountAssetsByIDs: jest.fn(),
       getNativeAssetTypes: jest
         .fn()
         .mockReturnValue([KnownCaip19Id.SolMainnet]),
@@ -143,7 +144,7 @@ describe('SolanaKeyring', () => {
   describe('getAccountAssets', () => {
     it('calls the assets service', async () => {
       jest
-        .spyOn(mockAssetsService, 'findByAccount')
+        .spyOn(mockAssetsService, 'getAccountAssetsForAllActiveScopes')
         .mockResolvedValue(MOCK_ASSET_ENTITIES);
 
       const result = await keyring.getAccountAssets(
@@ -158,10 +159,12 @@ describe('SolanaKeyring', () => {
     });
 
     it('removes token assets with zero balance', async () => {
-      jest.spyOn(mockAssetsService, 'findByAccount').mockResolvedValue([
-        MOCK_ASSET_ENTITY_1, // Token asset with non-zero balance
-        { ...MOCK_ASSET_ENTITY_2, rawAmount: '0' }, // Token asset with zero balance
-      ]);
+      jest
+        .spyOn(mockAssetsService, 'getAccountAssetsForAllActiveScopes')
+        .mockResolvedValue([
+          MOCK_ASSET_ENTITY_1, // Token asset with non-zero balance
+          { ...MOCK_ASSET_ENTITY_2, rawAmount: '0' }, // Token asset with zero balance
+        ]);
 
       const result = await keyring.getAccountAssets(
         MOCK_SOLANA_KEYRING_ACCOUNT_0.id,
@@ -171,10 +174,12 @@ describe('SolanaKeyring', () => {
     });
 
     it('keeps the native asset even if it has zero balance', async () => {
-      jest.spyOn(mockAssetsService, 'findByAccount').mockResolvedValue([
-        { ...MOCK_ASSET_ENTITY_0, rawAmount: '0' }, // Native asset with zero balance
-        { ...MOCK_ASSET_ENTITY_1, rawAmount: '0' }, // Token asset with zero balance
-      ]);
+      jest
+        .spyOn(mockAssetsService, 'getAccountAssetsForAllActiveScopes')
+        .mockResolvedValue([
+          { ...MOCK_ASSET_ENTITY_0, rawAmount: '0' }, // Native asset with zero balance
+          { ...MOCK_ASSET_ENTITY_1, rawAmount: '0' }, // Token asset with zero balance
+        ]);
 
       const result = await keyring.getAccountAssets(
         MOCK_SOLANA_KEYRING_ACCOUNT_0.id,
@@ -343,9 +348,9 @@ describe('SolanaKeyring', () => {
         symbol: 4,
       } as unknown as AssetEntity;
 
-      jest
-        .spyOn(mockAssetsService, 'findByAccount')
-        .mockResolvedValue([invalidAsset]);
+      jest.spyOn(mockAssetsService, 'getAccountAssetsByIDs').mockResolvedValue({
+        [KnownCaip19Id.SolMainnet]: invalidAsset,
+      });
 
       await expect(
         keyring.getAccountBalances(MOCK_SOLANA_KEYRING_ACCOUNT_1.id, [
@@ -355,10 +360,13 @@ describe('SolanaKeyring', () => {
     });
 
     it('removes token assets with zero balance', async () => {
-      jest.spyOn(mockAssetsService, 'findByAccount').mockResolvedValue([
-        MOCK_ASSET_ENTITY_1, // Token asset with non-zero balance
-        { ...MOCK_ASSET_ENTITY_2, rawAmount: '0' }, // Token asset with zero balance
-      ]);
+      jest.spyOn(mockAssetsService, 'getAccountAssetsByIDs').mockResolvedValue({
+        [MOCK_ASSET_ENTITY_1.assetType]: MOCK_ASSET_ENTITY_1,
+        [MOCK_ASSET_ENTITY_2.assetType]: {
+          ...MOCK_ASSET_ENTITY_2,
+          rawAmount: '0',
+        },
+      });
 
       const result = await keyring.getAccountBalances(
         MOCK_SOLANA_KEYRING_ACCOUNT_0.id,
@@ -374,10 +382,16 @@ describe('SolanaKeyring', () => {
     });
 
     it('keeps the native asset even if it has zero balance', async () => {
-      jest.spyOn(mockAssetsService, 'findByAccount').mockResolvedValue([
-        { ...MOCK_ASSET_ENTITY_0, rawAmount: '0' }, // Native asset with zero balance
-        { ...MOCK_ASSET_ENTITY_1, rawAmount: '0' }, // Token asset with zero balance
-      ]);
+      jest.spyOn(mockAssetsService, 'getAccountAssetsByIDs').mockResolvedValue({
+        [MOCK_ASSET_ENTITY_0.assetType]: {
+          ...MOCK_ASSET_ENTITY_0,
+          rawAmount: '0',
+        },
+        [MOCK_ASSET_ENTITY_1.assetType]: {
+          ...MOCK_ASSET_ENTITY_1,
+          rawAmount: '0',
+        },
+      });
 
       const result = await keyring.getAccountBalances(
         MOCK_SOLANA_KEYRING_ACCOUNT_0.id,

packages/snap/src/core/handlers/onKeyringRequest/Keyring.ts

@@ -415,7 +415,8 @@ export class SolanaKeyring implements KeyringSnapRpc {
 
       const account = await this.getAccountOrThrow(accountId);
 
-      const assetEntities = await this.#assetsService.findByAccount(account);
+      const assetEntities =
+        await this.#assetsService.getAccountAssetsForAllActiveScopes(accountId);
 
       const result = assetEntities
         // Remove token assets with zero balance
@@ -450,8 +451,13 @@ export class SolanaKeyring implements KeyringSnapRpc {
 
       const account = await this.getAccountOrThrow(accountId);
 
-      const assetsToUse = (await this.#assetsService.findByAccount(account))
-        .filter((asset) => assets.includes(asset.assetType))
+      const assetsById = await this.#assetsService.getAccountAssetsByIDs(
+        accountId,
+        assets,
+      );
+
+      const assetsToUse = Object.values(assetsById)
+        .filter((asset): asset is NonNullable<typeof asset> => asset !== null)
         // Remove token assets with zero balance
         .filter(
           (asset) =>

packages/snap/src/core/services/accounts/AccountsSynchronizer.ts

@@ -33,16 +33,26 @@ export class AccountsSynchronizer {
 
     const assets = (
       await Promise.allSettled(
-        accountsToSync.map(async (account) =>
-          this.#assetsService.fetch(account),
-        ),
+        accountsToSync.map(async (account) => {
+          if (
+            await this.#assetsService.shouldTrackSnapAssetsForAccount(
+              account.id,
+            )
+          ) {
+            const fetchedAssets = await this.#assetsService.fetch(account);
+            await this.#assetsService.saveMany(fetchedAssets);
+            return fetchedAssets;
+          }
+
+          return this.#assetsService.getAccountAssetsForAllActiveScopes(
+            account.id,
+          );
+        }),
       )
     )
       .map((item) => (item.status === 'fulfilled' ? item.value : []))
       .flat();
 
-    await this.#assetsService.saveMany(assets);
-
     const transactions =
       await this.#transactionsService.fetchAssetsTransactions(assets, {
         limit: 20,