[FEATURE FLAG] feat: migrate assets domain to use the AssetsController

packages/tron-wallet-snap/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add feature-flagged AssetsController migration for fungible asset reads (TRX, TRC10, TRC20). Protocol assets (energy, bandwidth, staking, lock/withdrawal, rewards) remain Snap-owned. Resolution order: remote feature flags → `TRON_ASSETS_MIGRATION_STAGE` env (non-production only) → Off default.
+
 ### Changed
 
 - Bump `@metamask/keyring-api` from `^23.2.0` to `^23.7.0` ([#43](https://github.com/MetaMask/internal-snaps/pull/43))

packages/tron-wallet-snap/package.json

@@ -48,10 +48,13 @@
     "test:watch": "NODE_OPTIONS=--experimental-vm-modules jest --watch"
   },
   "devDependencies": {
+    "@metamask/assets-controller": "11.2.0",
     "@metamask/auto-changelog": "^6.1.1",
     "@metamask/key-tree": "^10.1.1",
     "@metamask/keyring-api": "^23.7.0",
     "@metamask/keyring-snap-sdk": "^9.2.1",
+    "@metamask/messenger": "^2.0.0",
+    "@metamask/remote-feature-flag-controller": "4.2.2",
     "@metamask/snaps-cli": "^8.4.1",
     "@metamask/snaps-jest": "^10.2.0",
     "@metamask/snaps-sdk": "^11.2.0",

packages/tron-wallet-snap/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/internal-snaps.git"
   },
   "source": {
-    "shasum": "tcnp/6xxntfergEzeQYc7k96SfveVo+5kBQHGEXqZkQ=",
+    "shasum": "8FiFygUypvBGiZUEI1rAWWUTtysoQ7d4SHSbYmpYnnk=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",
@@ -48,6 +48,12 @@
     },
     "endowment:assets": {
       "scopes": ["tron:728126428"]
+    },
+    "endowment:messenger": {
+      "actions": [
+        "AssetsController:getAsset",
+        "RemoteFeatureFlagController:getState"
+      ]
     }
   },
   "platformVersion": "11.2.0",

packages/tron-wallet-snap/src/clients/price-api/PriceApiClient.ts

@@ -7,7 +7,7 @@ import { mapKeys } from 'lodash';
 
 import type { ICache } from '../../caching/ICache';
 import { useCache } from '../../caching/useCache';
-import { SPECIAL_ASSETS } from '../../constants';
+import { SNAP_OWNED_ASSETS } from '../../constants';
 import type { ConfigProvider } from '../../services/config';
 import { buildUrl } from '../../utils/buildUrl';
 import type { ILogger } from '../../utils/logger';
@@ -257,7 +257,7 @@ export class PriceApiClient {
     assert(vsCurrency, VsCurrencyParamStruct);
 
     const filteredTokens = tokenCaip19Types.filter(
-      (tokenCaip19Type) => !SPECIAL_ASSETS.includes(tokenCaip19Type),
+      (tokenCaip19Type) => !SNAP_OWNED_ASSETS.includes(tokenCaip19Type),
     );
 
     return this.#getMultipleSpotPrices_CACHE(filteredTokens, vsCurrency);

packages/tron-wallet-snap/src/clients/token-api/TokenApiClient.ts

@@ -3,7 +3,7 @@ import { array, assert } from '@metamask/superstruct';
 import type { Infer } from '@metamask/superstruct';
 import { CaipAssetTypeStruct, parseCaipAssetType } from '@metamask/utils';
 
-import { Network, SPECIAL_ASSETS } from '../../constants';
+import { Network, SNAP_OWNED_ASSETS } from '../../constants';
 import type { TokenCaipAssetType } from '../../services/assets/types';
 import { TokenCaipAssetTypeStruct } from '../../services/assets/types';
 import type { ConfigProvider } from '../../services/config';
@@ -94,7 +94,7 @@ export class TokenApiClient {
        * Exclude TRON resource tokens (energy and bandwidth), staked tokens, and tokens not from supported networks.
        */
       const supportedAssetTypes = assetTypes.filter((assetType) => {
-        if (SPECIAL_ASSETS.includes(assetType)) {
+        if (SNAP_OWNED_ASSETS.includes(assetType)) {
           return false;
         }
         const { chainId } = parseCaipAssetType(assetType);

packages/tron-wallet-snap/src/constants/index.ts

@@ -389,7 +389,7 @@ export const Networks = {
   },
 } as const;
 
-export const SPECIAL_ASSETS: string[] = [
+export const SNAP_OWNED_ASSETS: string[] = [
   KnownCaip19Id.TrxStakedForBandwidthMainnet,
   KnownCaip19Id.TrxStakedForBandwidthNile,
   KnownCaip19Id.TrxStakedForBandwidthShasta,
@@ -419,9 +419,12 @@ export const SPECIAL_ASSETS: string[] = [
   KnownCaip19Id.MaximumEnergyShasta,
 ];
 
+/** @deprecated Use {@link SNAP_OWNED_ASSETS} instead. */
+export const SPECIAL_ASSETS = SNAP_OWNED_ASSETS;
+
 export const ESSENTIAL_ASSETS: string[] = [
   KnownCaip19Id.TrxMainnet,
   KnownCaip19Id.TrxNile,
   KnownCaip19Id.TrxShasta,
-  ...SPECIAL_ASSETS,
+  ...SNAP_OWNED_ASSETS,
 ];

packages/tron-wallet-snap/src/context.ts

@@ -1,3 +1,5 @@
+import { getMessenger } from '@metamask/snaps-sdk';
+
 import { InMemoryCache } from './caching/InMemoryCache';
 import { StateCache } from './caching/StateCache';
 import { PriceApiClient } from './clients/price-api/PriceApiClient';
@@ -29,6 +31,7 @@ import { TransactionScanService } from './services/transaction-scan/TransactionS
 import { TransactionsRepository } from './services/transactions/TransactionsRepository';
 import { TransactionsService } from './services/transactions/TransactionsService';
 import { WalletService } from './services/wallet/WalletService';
+import type { CoreMessenger } from './types/core-messenger';
 import logger, { noOpLogger } from './utils/logger';
 
 /**
@@ -82,6 +85,8 @@ const priceApiClient = new PriceApiClient(configProvider, priceCache);
 // Token API client
 const tokenApiClient = new TokenApiClient(configProvider);
 
+const coreMessenger = getMessenger<CoreMessenger>();
+
 // Security Alerts API client
 const securityAlertsApiClient = new SecurityAlertsApiClient(
   configProvider,
@@ -98,6 +103,7 @@ const assetsService = new AssetsService({
   priceApiClient,
   tokenApiClient,
   snapClient,
+  coreMessenger,
 });
 
 const transactionsService = new TransactionsService({

packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.ts

@@ -369,13 +369,16 @@ export class ClientRequestHandler {
       const { chainId } = parseCaipAssetType(assetId);
       const scope = chainId as Network;
 
-      const [asset, nativeTokenAsset, bandwidthAsset, energyAsset] =
-        await this.#assetsService.getAssetsByAccountId(accountId, [
-          assetId,
-          Networks[scope].nativeToken.id,
-          Networks[scope].bandwidth.id,
-          Networks[scope].energy.id,
-        ]);
+      const assets = await this.#assetsService.getAccountAssetsByIDs(accountId, [
+        assetId,
+        Networks[scope].nativeToken.id,
+        Networks[scope].bandwidth.id,
+        Networks[scope].energy.id,
+      ]);
+      const asset = assets[assetId];
+      const nativeTokenAsset = assets[Networks[scope].nativeToken.id];
+      const bandwidthAsset = assets[Networks[scope].bandwidth.id];
+      const energyAsset = assets[Networks[scope].energy.id];
 
       const valueBN = new BigNumber(value);
       const assetToSendBalance = asset ? new BigNumber(asset.uiAmount) : ZERO;
@@ -483,7 +486,7 @@ export class ClientRequestHandler {
       };
     }
 
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       fromAccountId,
       assetId,
     );
@@ -523,11 +526,11 @@ export class ClientRequestHandler {
       };
     }
 
-    const [[bandwidthAsset, energyAsset], transaction] = await Promise.all([
+    const [bandwidthAndEnergyAssets, transaction] = await Promise.all([
       /**
        * Get available Energy and Bandwidth from account assets.
        */
-      this.#assetsService.getAssetsByAccountId(fromAccountId, [
+      this.#assetsService.getAccountAssetsByIDs(fromAccountId, [
         Networks[scope].bandwidth.id,
         Networks[scope].energy.id,
       ]),
@@ -545,6 +548,10 @@ export class ClientRequestHandler {
       }),
     ]);
 
+    const bandwidthAsset =
+      bandwidthAndEnergyAssets[Networks[scope].bandwidth.id];
+    const energyAsset = bandwidthAndEnergyAssets[Networks[scope].energy.id];
+
     const availableEnergy = energyAsset
       ? new BigNumber(energyAsset.rawAmount)
       : ZERO;
@@ -672,11 +679,14 @@ export class ClientRequestHandler {
     /**
      * Get available Energy and Bandwidth from account assets.
      */
-    const [bandwidthAsset, energyAsset] =
-      await this.#assetsService.getAssetsByAccountId(accountId, [
+    const bandwidthAndEnergyAssets =
+      await this.#assetsService.getAccountAssetsByIDs(accountId, [
         Networks[scope].bandwidth.id,
         Networks[scope].energy.id,
       ]);
+    const bandwidthAsset =
+      bandwidthAndEnergyAssets[Networks[scope].bandwidth.id];
+    const energyAsset = bandwidthAndEnergyAssets[Networks[scope].energy.id];
 
     const availableEnergy = energyAsset
       ? new BigNumber(energyAsset.rawAmount)
@@ -728,7 +738,7 @@ export class ClientRequestHandler {
 
     const scope = Network.Mainnet;
 
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       fromAccountId,
       Networks[scope].nativeToken.id,
     );
@@ -758,11 +768,14 @@ export class ClientRequestHandler {
     /**
      * Get available Energy and Bandwidth from account assets.
      */
-    const [bandwidthAsset, energyAsset] =
-      await this.#assetsService.getAssetsByAccountId(fromAccountId, [
+    const bandwidthAndEnergyAssets =
+      await this.#assetsService.getAccountAssetsByIDs(fromAccountId, [
         Networks[scope].bandwidth.id,
         Networks[scope].energy.id,
       ]);
+    const bandwidthAsset =
+      bandwidthAndEnergyAssets[Networks[scope].bandwidth.id];
+    const energyAsset = bandwidthAndEnergyAssets[Networks[scope].energy.id];
 
     const availableEnergy = energyAsset
       ? BigNumber(energyAsset.rawAmount)
@@ -804,7 +817,7 @@ export class ClientRequestHandler {
     const { accountId, assetId, value } = request.params;
 
     await this.#accountsService.findByIdOrThrow(accountId);
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       accountId,
       assetId,
     );
@@ -851,7 +864,7 @@ export class ClientRequestHandler {
 
     const account = await this.#accountsService.findByIdOrThrow(fromAccountId);
 
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       fromAccountId,
       assetId,
     );
@@ -913,7 +926,7 @@ export class ClientRequestHandler {
     const stakedAssetId = `${assetId}-staked-for-${purpose.toLowerCase()}`;
 
     await this.#accountsService.findByIdOrThrow(accountId);
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       accountId,
       stakedAssetId,
     );
@@ -967,7 +980,7 @@ export class ClientRequestHandler {
 
     const account = await this.#accountsService.findByIdOrThrow(accountId);
 
-    const asset = await this.#assetsService.getAssetByAccountId(
+    const asset = await this.#assetsService.getAccountAssetByID(
       accountId,
       stakedAssetId,
     );

packages/tron-wallet-snap/src/handlers/keyring.ts

@@ -196,12 +196,20 @@ export class KeyringHandler implements Keyring {
     try {
       validateRequest({ accountId }, ListAccountAssetsStruct);
 
-      await this.#getAccountOrThrow(accountId);
+      const account = await this.#getAccountOrThrow(accountId);
 
       this.#logger.info('Listing account assets', { accountId });
 
-      const assetEntities =
-        await this.#assetsService.getByKeyringAccountId(accountId);
+      const assetEntities = (
+        await Promise.all(
+          account.scopes.map((scope) =>
+            this.#assetsService.getAccountAssetsByScope(
+              scope as Network,
+              accountId,
+            ),
+          ),
+        )
+      ).flat();
       const result = assetEntities
         .filter(
           (asset) =>
@@ -342,10 +350,18 @@ export class KeyringHandler implements Keyring {
 
       this.#logger.info('Getting account balances', { accountId, assets });
 
-      await this.#getAccountOrThrow(accountId);
+      const account = await this.#getAccountOrThrow(accountId);
 
-      const assetsList =
-        await this.#assetsService.getByKeyringAccountId(accountId);
+      const assetsList = (
+        await Promise.all(
+          account.scopes.map((scope) =>
+            this.#assetsService.getAccountAssetsByScope(
+              scope as Network,
+              accountId,
+            ),
+          ),
+        )
+      ).flat();
 
       const assetsToUse = assetsList
         .filter((asset) => assets.includes(asset.assetType))