feat: add flag to exclude networks from using infura for live balance

packages/transaction-pay-controller/CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Live token balance queries now respect the `confirmations_pay_extended.excludeChainIdsFromInfura` feature flag, skipping the Infura endpoint preference for excluded chains ([#8992](https://github.com/MetaMask/core/pull/8992))
 - Bump `@metamask/assets-controllers` from `^108.3.0` to `^108.4.0` ([#8981](https://github.com/MetaMask/core/pull/8981))
 - Bump `@metamask/assets-controller` from `^8.0.2` to `^8.3.1` ([#8981](https://github.com/MetaMask/core/pull/8981), [#8985](https://github.com/MetaMask/core/pull/8985))
 - Bump `@metamask/remote-feature-flag-controller` from `^4.2.1` to `^4.2.2` ([#8986](https://github.com/MetaMask/core/pull/8986))

packages/transaction-pay-controller/src/utils/feature-flags.test.ts

@@ -23,6 +23,7 @@ import {
   getRelayOriginGasOverhead,
   getRelayPollingInterval,
   getRelayPollingTimeout,
+  isChainExcludedFromInfura,
   isEIP7702Chain,
   isRelayExecuteEnabled,
   getFeatureFlags,
@@ -502,6 +503,64 @@ describe('Feature Flags Utils', () => {
     });
   });
 
+  describe('isChainExcludedFromInfura', () => {
+    it('returns false when no feature flags are set', () => {
+      expect(isChainExcludedFromInfura(messenger, CHAIN_ID_MOCK)).toBe(false);
+    });
+
+    it('returns false when excludeChainIdsFromInfura is empty', () => {
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: [],
+          },
+        },
+      });
+
+      expect(isChainExcludedFromInfura(messenger, CHAIN_ID_MOCK)).toBe(false);
+    });
+
+    it('returns true when chainId is in the exclusion list', () => {
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: [CHAIN_ID_MOCK],
+          },
+        },
+      });
+
+      expect(isChainExcludedFromInfura(messenger, CHAIN_ID_MOCK)).toBe(true);
+    });
+
+    it('returns false when chainId is not in the exclusion list', () => {
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: [CHAIN_ID_DIFFERENT_MOCK],
+          },
+        },
+      });
+
+      expect(isChainExcludedFromInfura(messenger, CHAIN_ID_MOCK)).toBe(false);
+    });
+
+    it('performs case-insensitive comparison', () => {
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: ['0xA' as Hex],
+          },
+        },
+      });
+
+      expect(isChainExcludedFromInfura(messenger, '0xa' as Hex)).toBe(true);
+    });
+  });
+
   describe('getRelayOriginGasOverhead', () => {
     it('returns default when no feature flags are set', () => {
       expect(getRelayOriginGasOverhead(messenger)).toBe(

packages/transaction-pay-controller/src/utils/feature-flags.ts

@@ -149,6 +149,7 @@ export type PayStrategiesConfigRaw = {
 };
 
 type FeatureFlagsExtendedRaw = {
+  excludeChainIdsFromInfura?: Hex[];
   payStrategies?: {
     relay?: {
       gaslessEnabled?: boolean;
@@ -556,6 +557,33 @@ export function isRelayExecuteEnabled(
   return featureFlags.payStrategies?.relay?.gaslessEnabled ?? false;
 }
 
+/**
+ * Whether a chain is excluded from preferring Infura for balance queries.
+ *
+ * When a chain ID appears in the `confirmations_pay_extended.excludeChainIdsFromInfura`
+ * feature flag array, the Infura RPC endpoint should not be forced for that chain.
+ *
+ * @param messenger - Controller messenger.
+ * @param chainId - Chain ID to check.
+ * @returns True if the chain should skip the Infura preference.
+ */
+export function isChainExcludedFromInfura(
+  messenger: TransactionPayControllerMessenger,
+  chainId: Hex,
+): boolean {
+  const state = messenger.call('RemoteFeatureFlagController:getState');
+  const featureFlags =
+    (state.remoteFeatureFlags?.confirmations_pay_extended as
+      | FeatureFlagsExtendedRaw
+      | undefined) ?? {};
+
+  const excludedChains = featureFlags.excludeChainIdsFromInfura ?? [];
+
+  return excludedChains.some(
+    (excluded) => excluded.toLowerCase() === chainId.toLowerCase(),
+  );
+}
+
 /**
  * Get the origin gas overhead to include in Relay quote requests
  * for EIP-7702 chains.

packages/transaction-pay-controller/src/utils/token.test.ts

@@ -773,6 +773,79 @@ describe('Token Utils', () => {
         NETWORK_CLIENT_ID_MOCK,
       );
     });
+
+    it('skips Infura when chain is in excludeChainIdsFromInfura flag', async () => {
+      PROVIDER_MOCK.request.mockResolvedValue('0x4C4B40');
+
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: [CHAIN_ID_MOCK],
+          },
+        },
+      });
+
+      getNetworkConfigurationByChainIdMock.mockReturnValue({
+        rpcEndpoints: [
+          {
+            type: RpcEndpointType.Infura,
+            networkClientId: INFURA_NETWORK_CLIENT_ID_MOCK,
+          },
+        ],
+      } as NetworkConfiguration);
+
+      const result = await getLiveTokenBalance(
+        messenger,
+        ACCOUNT_MOCK,
+        CHAIN_ID_MOCK,
+        ERC20_ADDRESS_MOCK,
+      );
+
+      expect(result).toBe('5000000');
+      expect(getNetworkConfigurationByChainIdMock).not.toHaveBeenCalled();
+      expect(findNetworkClientIdByChainIdMock).toHaveBeenCalledWith(
+        CHAIN_ID_MOCK,
+      );
+      expect(getNetworkClientByIdMock).toHaveBeenCalledWith(
+        NETWORK_CLIENT_ID_MOCK,
+      );
+    });
+
+    it('uses Infura when chain is not in excludeChainIdsFromInfura flag', async () => {
+      PROVIDER_MOCK.request.mockResolvedValue('0x895440');
+
+      getRemoteFeatureFlagControllerStateMock.mockReturnValue({
+        ...getDefaultRemoteFeatureFlagControllerState(),
+        remoteFeatureFlags: {
+          confirmations_pay_extended: {
+            excludeChainIdsFromInfura: ['0x89' as Hex],
+          },
+        },
+      });
+
+      getNetworkConfigurationByChainIdMock.mockReturnValue({
+        rpcEndpoints: [
+          {
+            type: RpcEndpointType.Infura,
+            networkClientId: INFURA_NETWORK_CLIENT_ID_MOCK,
+          },
+        ],
+      } as NetworkConfiguration);
+
+      const result = await getLiveTokenBalance(
+        messenger,
+        ACCOUNT_MOCK,
+        CHAIN_ID_MOCK,
+        ERC20_ADDRESS_MOCK,
+      );
+
+      expect(result).toBe('9000000');
+      expect(getNetworkClientByIdMock).toHaveBeenCalledWith(
+        INFURA_NETWORK_CLIENT_ID_MOCK,
+      );
+      expect(findNetworkClientIdByChainIdMock).not.toHaveBeenCalled();
+    });
   });
 
   describe('computeTokenAmounts', () => {

packages/transaction-pay-controller/src/utils/token.ts

@@ -13,7 +13,10 @@ import {
   STABLECOINS,
 } from '../constants';
 import type { FiatRates, TransactionPayControllerMessenger } from '../types';
-import { getAssetsUnifyStateFeature } from './feature-flags';
+import {
+  getAssetsUnifyStateFeature,
+  isChainExcludedFromInfura,
+} from './feature-flags';
 import { getNetworkClientId, rpcRequest } from './provider';
 
 /**
@@ -322,7 +325,9 @@ export async function getLiveTokenBalance(
   chainId: Hex,
   tokenAddress: Hex,
 ): Promise<string> {
-  const options = { preferInfura: true };
+  const options = {
+    preferInfura: !isChainExcludedFromInfura(messenger, chainId),
+  };
   const isNative =
     tokenAddress.toLowerCase() === getNativeToken(chainId).toLowerCase();