feat: assets controller resilience

packages/assets-controller/src/AssetsController.test.ts

@@ -137,6 +137,11 @@ async function withController<ReturnValue>(
   ]: [WithControllerOptions, WithControllerCallback<ReturnValue>] =
     args.length === 2 ? args : [{}, args[0]];
 
+  const {
+    priceDataSourceConfig: incomingPriceDataSourceConfig,
+    ...restControllerOptions
+  } = controllerOptions;
+
   // Use root messenger (MOCK_ANY_NAMESPACE) so data sources can register their actions.
   const messenger: RootMessenger = new Messenger({
     namespace: MOCK_ANY_NAMESPACE,
@@ -207,7 +212,11 @@ async function withController<ReturnValue>(
     subscribeToBasicFunctionalityChange: (): void => {
       /* no-op for tests */
     },
-    ...controllerOptions,
+    ...restControllerOptions,
+    priceDataSourceConfig: {
+      simulateMiddlewareFailure: false,
+      ...incomingPriceDataSourceConfig,
+    },
   });
 
   try {
@@ -302,6 +311,7 @@ describe('AssetsController', () => {
         subscribeToBasicFunctionalityChange: (): void => {
           /* no-op for tests */
         },
+        priceDataSourceConfig: { simulateMiddlewareFailure: false },
       });
 
       // Controller should still have default state (from super() call)
@@ -360,6 +370,7 @@ describe('AssetsController', () => {
             subscribeToBasicFunctionalityChange: (): void => {
               /* no-op */
             },
+            priceDataSourceConfig: { simulateMiddlewareFailure: false },
             accountsApiDataSourceConfig: {
               pollInterval: 15_000,
               tokenDetectionEnabled: (): boolean => false,
@@ -400,6 +411,7 @@ describe('AssetsController', () => {
             },
             priceDataSourceConfig: {
               pollInterval: 120_000,
+              simulateMiddlewareFailure: false,
             },
           }),
       ).not.toThrow();
@@ -518,6 +530,99 @@ describe('AssetsController', () => {
     });
   });
 
+  describe('custom asset graduation', () => {
+    const SOLANA_ASSET_ID =
+      'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v' as Caip19AssetId;
+
+    it('graduates an EVM custom asset when AccountsApiDataSource reports a balance for it', async () => {
+      await withController(async ({ controller }) => {
+        await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID);
+        expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain(
+          MOCK_ASSET_ID,
+        );
+
+        await controller.handleAssetsUpdate(
+          {
+            assetsBalance: {
+              [MOCK_ACCOUNT_ID]: {
+                [MOCK_ASSET_ID]: { amount: '1000000' },
+              },
+            },
+          },
+          'AccountsApiDataSource',
+        );
+
+        expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined();
+      });
+    });
+
+    it('graduates an EVM custom asset when BackendWebsocketDataSource reports a balance for it', async () => {
+      await withController(async ({ controller }) => {
+        await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID);
+
+        await controller.handleAssetsUpdate(
+          {
+            assetsBalance: {
+              [MOCK_ACCOUNT_ID]: {
+                [MOCK_ASSET_ID]: { amount: '1000000' },
+              },
+            },
+          },
+          'BackendWebsocketDataSource',
+        );
+
+        expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toBeUndefined();
+      });
+    });
+
+    it('does not graduate when RpcDataSource reports a balance for a custom asset', async () => {
+      await withController(async ({ controller }) => {
+        await controller.addCustomAsset(MOCK_ACCOUNT_ID, MOCK_ASSET_ID);
+
+        await controller.handleAssetsUpdate(
+          {
+            assetsBalance: {
+              [MOCK_ACCOUNT_ID]: {
+                [MOCK_ASSET_ID]: { amount: '1000000' },
+              },
+            },
+          },
+          'RpcDataSource',
+        );
+
+        expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain(
+          MOCK_ASSET_ID,
+        );
+      });
+    });
+
+    it('does not graduate a non-EVM (Solana) custom asset', async () => {
+      await withController(
+        {
+          state: {
+            customAssets: { [MOCK_ACCOUNT_ID]: [SOLANA_ASSET_ID] },
+          },
+        },
+        async ({ controller }) => {
+          await controller.handleAssetsUpdate(
+            {
+              assetsBalance: {
+                [MOCK_ACCOUNT_ID]: {
+                  [SOLANA_ASSET_ID]: { amount: '1000000' },
+                },
+              },
+            },
+            'AccountsApiDataSource',
+          );
+
+          expect(controller.state.customAssets[MOCK_ACCOUNT_ID]).toContain(
+            SOLANA_ASSET_ID,
+          );
+        },
+      );
+    });
+  });
+
   describe('getCustomAssets', () => {
     it('returns empty array for account with no custom assets', async () => {
       await withController(({ controller }) => {
@@ -1849,6 +1954,7 @@ describe('AssetsController', () => {
         subscribeToBasicFunctionalityChange: (): void => {
           /* no-op */
         },
+        priceDataSourceConfig: { simulateMiddlewareFailure: false },
       });
 
       const getAssetsSpy = jest.spyOn(controller, 'getAssets');

packages/assets-controller/src/AssetsController.ts

@@ -81,7 +81,10 @@ import type { StakedBalanceDataSourceConfig } from './data-sources/StakedBalance
 import { StakedBalanceDataSource } from './data-sources/StakedBalanceDataSource';
 import { TokenDataSource } from './data-sources/TokenDataSource';
 import { projectLogger, createModuleLogger } from './logger';
+import { AssetsDataSourceError } from './errors';
+import { CustomAssetGraduationMiddleware } from './middlewares/CustomAssetGraduationMiddleware';
 import { DetectionMiddleware } from './middlewares/DetectionMiddleware';
+import { RpcFallbackMiddleware } from './middlewares/RpcFallbackMiddleware';
 import {
   createParallelBalanceMiddleware,
   createParallelMiddleware,
@@ -354,6 +357,13 @@ export type AssetsControllerOptions = {
    * Use this to report first init fetch duration to Sentry (e.g. via addBreadcrumb or setMeasurement).
    */
   trace?: TraceCallback;
+  /**
+   * Optional Sentry (or compatible) reporter for **issue** events. When data source middlewares
+   * fail, the controller constructs an `AssetsDataSourceError` and passes it here so you can call
+   * `Sentry.captureException`. Span-only trace callbacks do not create Issues;
+   * wire this if you need Sentry alerts on middleware failures.
+   */
+  captureException?: (error: Error) => void;
   /** Optional configuration for AccountsApiDataSource. */
   accountsApiDataSourceConfig?: AccountsApiDataSourceConfig;
   /** Optional configuration for PriceDataSource. */
@@ -533,6 +543,9 @@ export class AssetsController extends BaseController<
   /** Optional trace callback for first init/fetch measurement (duration). */
   readonly #trace?: TraceCallback;
 
+  /** Optional reporter for Issue-style errors (e.g. Sentry.captureException). */
+  readonly #captureException?: (error: Error) => void;
+
   /** Whether we have already reported first init fetch for this session (reset on #stop). */
   #firstInitFetchReported = false;
 
@@ -697,6 +710,10 @@ export class AssetsController extends BaseController<
 
   readonly #detectionMiddleware: DetectionMiddleware;
 
+  readonly #customAssetGraduationMiddleware: CustomAssetGraduationMiddleware;
+
+  readonly #rpcFallbackMiddleware: RpcFallbackMiddleware;
+
   readonly #tokenDataSource: TokenDataSource;
 
   #unsubscribeBasicFunctionality: (() => void) | null = null;
@@ -717,6 +734,7 @@ export class AssetsController extends BaseController<
     queryApiClient,
     rpcDataSourceConfig,
     trace,
+    captureException,
     accountsApiDataSourceConfig,
     priceDataSourceConfig,
     stakedBalanceDataSourceConfig,
@@ -736,6 +754,7 @@ export class AssetsController extends BaseController<
     this.#isBasicFunctionality = isBasicFunctionality ?? ((): boolean => true);
     this.#defaultUpdateInterval = defaultUpdateInterval;
     this.#trace = trace;
+    this.#captureException = captureException;
     const rpcConfig = rpcDataSourceConfig ?? {};
 
     this.#onActiveChainsUpdated = (
@@ -791,8 +810,26 @@ export class AssetsController extends BaseController<
       queryApiClient,
       getSelectedCurrency: (): SupportedCurrency => this.state.selectedCurrency,
       ...priceDataSourceConfig,
+      simulateMiddlewareFailure:
+        priceDataSourceConfig?.simulateMiddlewareFailure ?? true,
     });
     this.#detectionMiddleware = new DetectionMiddleware();
+    this.#customAssetGraduationMiddleware = new CustomAssetGraduationMiddleware(
+      {
+        getSelectedAccountId: (): AccountId | undefined => {
+          try {
+            return this.#getSelectedAccounts()[0]?.id;
+          } catch {
+            return undefined;
+          }
+        },
+        removeCustomAsset: (accountId, assetId): void =>
+          this.removeCustomAsset(accountId, assetId),
+      },
+    );
+    this.#rpcFallbackMiddleware = new RpcFallbackMiddleware({
+      rpcDataSource: this.#rpcDataSource,
+    });
 
     if (!this.#isEnabled) {
       log('AssetsController is disabled, skipping initialization');
@@ -1212,13 +1249,32 @@ export class AssetsController extends BaseController<
       });
     }
 
-    // Emit error traces for failed middlewares
+    // Failed middlewares: Issues (optional) + perf/Dashboard spans
     if (middlewareErrors.length > 0) {
-      this.#emitTrace(TRACE_DATA_SOURCE_ERROR, {
-        failed_sources: middlewareErrors.join(','),
-        error_count: middlewareErrors.length,
-        chain_count: request.chainIds.length,
+      const failedSources = middlewareErrors.join(',');
+      const assetsError = new AssetsDataSourceError({
+        failedSources,
+        errorCount: middlewareErrors.length,
+        chainCount: request.chainIds.length,
       });
+      try {
+        this.#captureException?.(assetsError);
+      } catch {
+        // Never let telemetry throw.
+      }
+      this.#emitTrace(
+        TRACE_DATA_SOURCE_ERROR,
+        {
+          failed_sources: failedSources,
+          error_count: middlewareErrors.length,
+          chain_count: request.chainIds.length,
+        },
+        {
+          controller: 'AssetsController',
+          severity: 'error',
+          error_type: assetsError.name,
+        },
+      );
     }
 
     return { response: result.response, durationByDataSource };
@@ -1278,6 +1334,8 @@ export class AssetsController extends BaseController<
               this.#accountsApiDataSource,
               this.#stakedBalanceDataSource,
             ]),
+            this.#rpcFallbackMiddleware,
+            this.#customAssetGraduationMiddleware,
             this.#detectionMiddleware,
             createParallelMiddleware([
               this.#tokenDataSource,
@@ -2678,7 +2736,19 @@ export class AssetsController extends BaseController<
           ),
         };
 
-    const enrichmentSources: AssetsDataSource[] = [this.#detectionMiddleware];
+    // Graduate custom assets only when AccountsAPI / Websocket reports them.
+    // RPC already fetches custom assets on purpose, and Snap handles non-EVM
+    // chains the rule does not apply to, so skip the middleware for those.
+    const shouldGraduateCustomAssets =
+      sourceId === 'AccountsApiDataSource' ||
+      sourceId === 'BackendWebsocketDataSource';
+
+    const enrichmentSources: AssetsDataSource[] = [
+      ...(shouldGraduateCustomAssets
+        ? [this.#customAssetGraduationMiddleware]
+        : []),
+      this.#detectionMiddleware,
+    ];
     if (this.#isBasicFunctionality()) {
       enrichmentSources.push(
         createParallelMiddleware([

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

@@ -121,12 +121,14 @@ async function setupController(
     supportedChains?: number[];
     balances?: V5BalanceItem[];
     unprocessedNetworks?: string[];
+    fetchTimeoutMs?: number;
   } = {},
 ): Promise<SetupResult> {
   const {
     supportedChains = [1, 137],
     balances = [],
     unprocessedNetworks = [],
+    fetchTimeoutMs,
   } = options;
 
   const rootMessenger = new Messenger<MockAnyNamespace, AllActions, AllEvents>({
@@ -163,6 +165,7 @@ async function setupController(
       apiClient as unknown as AccountsApiDataSourceOptions['queryApiClient'],
     onActiveChainsUpdated: (dataSourceName, chains, previousChains): void =>
       activeChainsUpdateHandler(dataSourceName, chains, previousChains),
+    ...(fetchTimeoutMs === undefined ? {} : { fetchTimeoutMs }),
   });
 
   // Wait for async initialization
@@ -336,6 +339,24 @@ describe('AccountsApiDataSource', () => {
     controller.destroy();
   });
 
+  it('fetch marks every requested chain as errored when the call exceeds the configured timeout', async () => {
+    const { controller, apiClient } = await setupController({
+      fetchTimeoutMs: 10,
+    });
+
+    apiClient.accounts.fetchV5MultiAccountBalances.mockImplementationOnce(
+      () => new Promise(() => undefined),
+    );
+
+    const response = await controller.fetch(
+      createDataRequest({ chainIds: [CHAIN_MAINNET] }),
+    );
+
+    expect(response.errors?.[CHAIN_MAINNET]).toContain('timed out');
+
+    controller.destroy();
+  });
+
   it('fetch skips API when no valid account-chain combinations', async () => {
     const { controller, apiClient } = await setupController();