fix(transaction-pay-controller): use Infura endpoint for live token balance queries

packages/transaction-pay-controller/CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Bump `@metamask/gas-fee-controller` from `^26.2.1` to `^26.2.2` ([#8834](https://github.com/MetaMask/core/pull/8834))
+- `getLiveTokenBalance` now prefers the Infura RPC endpoint for a chain when querying live token balances, falling back to the chain's default endpoint if no Infura endpoint is configured ([#XXXX](https://github.com/MetaMask/core/pull/XXXX))
 
 ### Fixed
 

packages/transaction-pay-controller/src/strategy/across/across-submit.ts

@@ -21,6 +21,7 @@ import type {
 import { accountSupports7702 } from '../../utils/7702';
 import { getPayStrategiesConfig } from '../../utils/feature-flags';
 import { getGasBuffer } from '../../utils/feature-flags';
+import { getNetworkClientId } from '../../utils/provider';
 import {
   collectTransactionIds,
   getTransaction,
@@ -141,10 +142,7 @@ async function submitTransactions(
   const transactionCount =
     orderedTransactions.length + (shouldPrependOriginalTransaction ? 1 : 0);
 
-  const networkClientId = messenger.call(
-    'NetworkController:findNetworkClientIdByChainId',
-    chainId,
-  );
+  const networkClientId = getNetworkClientId(messenger, chainId);
 
   const is7702Batch = is7702 && transactionCount > 1;
   const canUseQuotedBatchGasLimit =

packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts

@@ -20,6 +20,7 @@ import {
   getRelayPollingInterval,
   getRelayPollingTimeout,
 } from '../../utils/feature-flags';
+import { getNetworkClientId } from '../../utils/provider';
 import {
   getLiveTokenBalance,
   normalizeTokenAddress,
@@ -488,10 +489,7 @@ async function submitViaRelayExecute(
   const { from, sourceChainId } = quote.request;
   const { requestId } = quote.original.steps[0];
 
-  const networkClientId = messenger.call(
-    'NetworkController:findNetworkClientIdByChainId',
-    sourceChainId,
-  );
+  const networkClientId = getNetworkClientId(messenger, sourceChainId);
 
   const sourceCallTransaction = {
     ...transaction,
@@ -580,10 +578,7 @@ async function submitViaTransactionController(
   const { from, sourceChainId, sourceTokenAddress } = quote.request;
   const { isPostQuote } = quote.request;
 
-  const networkClientId = messenger.call(
-    'NetworkController:findNetworkClientIdByChainId',
-    sourceChainId,
-  );
+  const networkClientId = getNetworkClientId(messenger, sourceChainId);
 
   log('Adding transactions', {
     normalizedParams: allParams,

packages/transaction-pay-controller/src/tests/messenger-mock.ts

@@ -15,6 +15,7 @@ import type {
 import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger';
 import type { NetworkControllerGetNetworkClientByIdAction } from '@metamask/network-controller';
 import type { NetworkControllerFindNetworkClientIdByChainIdAction } from '@metamask/network-controller';
+import type { NetworkControllerGetNetworkConfigurationByChainIdAction } from '@metamask/network-controller';
 import type { RemoteFeatureFlagControllerGetStateAction } from '@metamask/remote-feature-flag-controller';
 import type {
   TransactionControllerAddTransactionAction,
@@ -116,6 +117,10 @@ export function getMessengerMock({
     NetworkControllerGetNetworkClientByIdAction['handler']
   > = jest.fn();
 
+  const getNetworkConfigurationByChainIdMock: jest.MockedFn<
+    NetworkControllerGetNetworkConfigurationByChainIdAction['handler']
+  > = jest.fn();
+
   const getDelegationTransactionMock: jest.MockedFn<
     TransactionPayControllerGetDelegationTransactionAction['handler']
   > = jest.fn();
@@ -250,6 +255,11 @@ export function getMessengerMock({
       getNetworkClientByIdMock,
     );
 
+    messenger.registerActionHandler(
+      'NetworkController:getNetworkConfigurationByChainId',
+      getNetworkConfigurationByChainIdMock,
+    );
+
     messenger.registerActionHandler(
       'TransactionPayController:getDelegationTransaction',
       getDelegationTransactionMock,
@@ -310,6 +320,7 @@ export function getMessengerMock({
     getGasFeeTokensMock,
     getKeyringControllerStateMock,
     getNetworkClientByIdMock,
+    getNetworkConfigurationByChainIdMock,
     getRemoteFeatureFlagControllerStateMock,
     getStrategyMock,
     getTokenBalanceControllerStateMock,

packages/transaction-pay-controller/src/types.ts

@@ -31,6 +31,7 @@ import type {
 import type { Messenger } from '@metamask/messenger';
 import type { NetworkControllerFindNetworkClientIdByChainIdAction } from '@metamask/network-controller';
 import type { NetworkControllerGetNetworkClientByIdAction } from '@metamask/network-controller';
+import type { NetworkControllerGetNetworkConfigurationByChainIdAction } from '@metamask/network-controller';
 import type { Quote as RampsQuote } from '@metamask/ramps-controller';
 import type {
   RampsControllerGetOrderAction,
@@ -73,6 +74,7 @@ export type AllowedActions =
   | KeyringControllerSignTypedMessageAction
   | NetworkControllerFindNetworkClientIdByChainIdAction
   | NetworkControllerGetNetworkClientByIdAction
+  | NetworkControllerGetNetworkConfigurationByChainIdAction
   | RampsControllerGetOrderAction
   | RampsControllerGetQuotesAction
   | RampsControllerGetStateAction

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

@@ -11,6 +11,7 @@ import type { TransactionPayControllerMessenger } from '..';
 import { createModuleLogger, projectLogger } from '../logger';
 import type { Amount } from '../types';
 import { getFallbackGas, getGasBuffer } from './feature-flags';
+import { getNetworkClientId } from './provider';
 import { getNativeToken, getTokenBalance, getTokenFiatRate } from './token';
 
 const log = createModuleLogger(projectLogger, 'gas');
@@ -227,10 +228,7 @@ export async function estimateGasLimit({
   error?: unknown;
 }> {
   const gasBuffer = getGasBuffer(messenger, chainId);
-  const networkClientId = messenger.call(
-    'NetworkController:findNetworkClientIdByChainId',
-    chainId,
-  );
+  const networkClientId = getNetworkClientId(messenger, chainId);
 
   let estimateGasError: unknown;
   let simulationError: Error | undefined;

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

@@ -0,0 +1,186 @@
+import type { Provider } from '@metamask/network-controller';
+import { RpcEndpointType } from '@metamask/network-controller';
+import type { NetworkConfiguration } from '@metamask/network-controller';
+import type { Hex } from '@metamask/utils';
+
+import { getMessengerMock } from '../tests/messenger-mock';
+import { getNetworkClientId, rpcRequest } from './provider';
+
+const CHAIN_ID_MOCK = '0x1' as Hex;
+const DEFAULT_NETWORK_CLIENT_ID_MOCK = 'default-client-id';
+const INFURA_NETWORK_CLIENT_ID_MOCK = 'mainnet';
+const PROVIDER_MOCK = { request: jest.fn() } as unknown as Provider;
+
+describe('provider utils', () => {
+  const {
+    messenger,
+    findNetworkClientIdByChainIdMock,
+    getNetworkClientByIdMock,
+    getNetworkConfigurationByChainIdMock,
+  } = getMessengerMock();
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+
+    findNetworkClientIdByChainIdMock.mockReturnValue(
+      DEFAULT_NETWORK_CLIENT_ID_MOCK,
+    );
+
+    getNetworkClientByIdMock.mockReturnValue({
+      provider: PROVIDER_MOCK,
+    } as never);
+
+    getNetworkConfigurationByChainIdMock.mockReturnValue(undefined);
+  });
+
+  describe('getNetworkClientId', () => {
+    it('returns default network client ID when preferInfura is false', () => {
+      const result = getNetworkClientId(messenger, CHAIN_ID_MOCK);
+
+      expect(result).toBe(DEFAULT_NETWORK_CLIENT_ID_MOCK);
+      expect(findNetworkClientIdByChainIdMock).toHaveBeenCalledWith(
+        CHAIN_ID_MOCK,
+      );
+      expect(getNetworkConfigurationByChainIdMock).not.toHaveBeenCalled();
+    });
+
+    it('returns Infura network client ID when preferInfura is true and Infura endpoint exists', () => {
+      getNetworkConfigurationByChainIdMock.mockReturnValue({
+        rpcEndpoints: [
+          {
+            type: RpcEndpointType.Infura,
+            networkClientId: INFURA_NETWORK_CLIENT_ID_MOCK,
+          },
+        ],
+      } as NetworkConfiguration);
+
+      const result = getNetworkClientId(messenger, CHAIN_ID_MOCK, {
+        preferInfura: true,
+      });
+
+      expect(result).toBe(INFURA_NETWORK_CLIENT_ID_MOCK);
+      expect(findNetworkClientIdByChainIdMock).not.toHaveBeenCalled();
+    });
+
+    it('falls back to default network client ID when preferInfura is true but no Infura endpoint exists', () => {
+      getNetworkConfigurationByChainIdMock.mockReturnValue({
+        rpcEndpoints: [
+          {
+            type: RpcEndpointType.Custom,
+            networkClientId: 'custom-rpc-id',
+          },
+        ],
+      } as NetworkConfiguration);
+
+      const result = getNetworkClientId(messenger, CHAIN_ID_MOCK, {
+        preferInfura: true,
+      });
+
+      expect(result).toBe(DEFAULT_NETWORK_CLIENT_ID_MOCK);
+      expect(findNetworkClientIdByChainIdMock).toHaveBeenCalledWith(
+        CHAIN_ID_MOCK,
+      );
+    });
+
+    it('falls back to default network client ID when preferInfura is true but getNetworkConfigurationByChainId throws', () => {
+      getNetworkConfigurationByChainIdMock.mockImplementation(() => {
+        throw new Error('Configuration not found');
+      });
+
+      const result = getNetworkClientId(messenger, CHAIN_ID_MOCK, {
+        preferInfura: true,
+      });
+
+      expect(result).toBe(DEFAULT_NETWORK_CLIENT_ID_MOCK);
+      expect(findNetworkClientIdByChainIdMock).toHaveBeenCalledWith(
+        CHAIN_ID_MOCK,
+      );
+    });
+
+    it('falls back to default network client ID when preferInfura is true but network configuration is undefined', () => {
+      getNetworkConfigurationByChainIdMock.mockReturnValue(undefined);
+
+      const result = getNetworkClientId(messenger, CHAIN_ID_MOCK, {
+        preferInfura: true,
+      });
+
+      expect(result).toBe(DEFAULT_NETWORK_CLIENT_ID_MOCK);
+      expect(findNetworkClientIdByChainIdMock).toHaveBeenCalledWith(
+        CHAIN_ID_MOCK,
+      );
+    });
+  });
+
+  describe('rpcRequest', () => {
+    it('calls provider.request with method and params', async () => {
+      const requestMock = jest.fn().mockResolvedValue('0xabc');
+      getNetworkClientByIdMock.mockReturnValue({
+        provider: { request: requestMock },
+      } as never);
+
+      const result = await rpcRequest(messenger, CHAIN_ID_MOCK, 'eth_chainId', [
+        'latest',
+      ]);
+
+      expect(result).toBe('0xabc');
+      expect(requestMock).toHaveBeenCalledWith({
+        method: 'eth_chainId',
+        params: ['latest'],
+      });
+    });
+
+    it('calls provider.request without params when omitted', async () => {
+      const requestMock = jest.fn().mockResolvedValue('0x10');
+      getNetworkClientByIdMock.mockReturnValue({
+        provider: { request: requestMock },
+      } as never);
+
+      await rpcRequest(messenger, CHAIN_ID_MOCK, 'eth_blockNumber');
+
+      expect(requestMock).toHaveBeenCalledWith({
+        method: 'eth_blockNumber',
+        params: undefined,
+      });
+    });
+
+    it('propagates provider errors', async () => {
+      const error = new Error('RPC failed');
+      const requestMock = jest.fn().mockRejectedValue(error);
+      getNetworkClientByIdMock.mockReturnValue({
+        provider: { request: requestMock },
+      } as never);
+
+      await expect(
+        rpcRequest(messenger, CHAIN_ID_MOCK, 'eth_blockNumber'),
+      ).rejects.toBe(error);
+    });
+
+    it('uses Infura network client when preferInfura is true', async () => {
+      getNetworkConfigurationByChainIdMock.mockReturnValue({
+        rpcEndpoints: [
+          {
+            type: RpcEndpointType.Infura,
+            networkClientId: INFURA_NETWORK_CLIENT_ID_MOCK,
+          },
+        ],
+      } as NetworkConfiguration);
+
+      const requestMock = jest.fn().mockResolvedValue('0x1');
+      getNetworkClientByIdMock.mockReturnValue({
+        provider: { request: requestMock },
+      } as never);
+
+      await rpcRequest(
+        messenger,
+        CHAIN_ID_MOCK,
+        'eth_chainId',
+        [],
+        { preferInfura: true },
+      );
+
+      expect(getNetworkClientByIdMock).toHaveBeenCalledWith(
+        INFURA_NETWORK_CLIENT_ID_MOCK,
+      );
+    });
+  });
+});