feat: extend isAtomicBatchSupported result

packages/transaction-controller/CHANGELOG.md

@@ -7,8 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add types for `isAtomicBatchSupported` method ([#5600](https://github.com/MetaMask/core/pull/5600))
+  - `IsAtomicBatchSupportedRequest`
+  - `IsAtomicBatchSupportedResult`
+  - `IsAtomicBatchSupportedResultEntry`
+
 ### Changed
 
+- **BREAKING:** Update signature of `isAtomicBatchSupported` method ([#5600](https://github.com/MetaMask/core/pull/5600))
+  - Replace `address` argument with `request` object containing `address` and optional `chainIds`.
+  - Return array of `IsAtomicBatchSupportedResultEntry` objects.
 - Skip `origin` validation for `batch` transaction type ([#5586](https://github.com/MetaMask/core/pull/5586))
 - **BREAKING:** `enableTxParamsGasFeeUpdates` is renamed to `isAutomaticGasFeeUpdateEnabled` now expects a callback function instead of a boolean.
   - This callback is invoked before performing `txParams` gas fee updates. The update will proceed only if the callback returns a truthy value.

packages/transaction-controller/src/TransactionController.ts

@@ -103,6 +103,8 @@ import type {
   PublishHook,
   PublishBatchHook,
   GasFeeToken,
+  IsAtomicBatchSupportedResult,
+  IsAtomicBatchSupportedRequest,
 } from './types';
 import {
   TransactionEnvelopeType,
@@ -1057,12 +1059,14 @@ export class TransactionController extends BaseController<
   /**
    * Determine which chains support atomic batch transactions with the given account address.
    *
-   * @param address - The address of the account to check.
-   * @returns  The supported chain IDs.
+   * @param request - Request object containing the account address and other parameters.
+   * @returns  Result object containing the supported chains and related information.
    */
-  async isAtomicBatchSupported(address: Hex): Promise<Hex[]> {
+  async isAtomicBatchSupported(
+    request: IsAtomicBatchSupportedRequest,
+  ): Promise<IsAtomicBatchSupportedResult> {
     return isAtomicBatchSupported({
-      address,
+      ...request,
       getEthQuery: (chainId) => this.#getEthQuery({ chainId }),
       messenger: this.messagingSystem,
       publicKeyEIP7702: this.#publicKeyEIP7702,

packages/transaction-controller/src/index.ts

@@ -44,6 +44,9 @@ export type {
   GasPriceGasFeeEstimates,
   GasPriceValue,
   InferTransactionTypeResult,
+  IsAtomicBatchSupportedRequest,
+  IsAtomicBatchSupportedResult,
+  IsAtomicBatchSupportedResultEntry,
   LegacyGasFeeEstimates,
   Log,
   NestedTransactionMetadata,

packages/transaction-controller/src/types.ts

@@ -1720,3 +1720,33 @@ export type GasFeeToken = {
   /** Address of the token contract. */
   tokenAddress: Hex;
 };
+
+/** Request to check if atomic batch is supported for an account. */
+export type IsAtomicBatchSupportedRequest = {
+  /** Address of the account to check. */
+  address: Hex;
+
+  /**
+   * IDs of specific chains to check.
+   * If not provided, all supported chains will be checked.
+   */
+  chainIds?: Hex[];
+};
+
+/** Result of checking if atomic batch is supported for an account. */
+export type IsAtomicBatchSupportedResult = IsAtomicBatchSupportedResultEntry[];
+
+/** Info about atomic batch support for a single chain. */
+export type IsAtomicBatchSupportedResultEntry = {
+  /** ID of the chain. */
+  chainId: Hex;
+
+  /** Address of the contract that the account was upgraded to. */
+  delegationAddress?: Hex;
+
+  /** Whether the upgraded contract is supported. */
+  isSupported: boolean;
+
+  /** Address of the contract that the account would be upgraded to. */
+  upgradeContractAddress?: Hex;
+};

packages/transaction-controller/src/utils/batch.test.ts

@@ -53,6 +53,8 @@ const TRANSACTION_SIGNATURE_MOCK = '0xabc';
 const TRANSACTION_SIGNATURE_2_MOCK = '0xdef';
 const ERROR_MESSAGE_MOCK = 'Test error';
 const SECURITY_ALERT_ID_MOCK = '123-456';
+const UPGRADE_CONTRACT_ADDRESS_MOCK =
+  '0xfedfedfedfedfedfedfedfedfedfedfedfedfedf';
 
 const TRANSACTION_META_MOCK = {
   id: BATCH_ID_CUSTOM_MOCK,
@@ -1073,12 +1075,16 @@ describe('Batch Utils', () => {
   });
 
   describe('isAtomicBatchSupported', () => {
-    it('includes feature flag chains if not upgraded or upgraded to supported contract', async () => {
+    it('includes all feature flag chains if chain IDs not specified', async () => {
       getEIP7702SupportedChainsMock.mockReturnValueOnce([
         CHAIN_ID_MOCK,
         CHAIN_ID_2_MOCK,
       ]);
 
+      getEIP7702UpgradeContractAddressMock.mockReturnValue(
+        UPGRADE_CONTRACT_ADDRESS_MOCK,
+      );
+
       isAccountUpgradedToEIP7702Mock
         .mockResolvedValueOnce({
           isSupported: false,
@@ -1096,25 +1102,53 @@ describe('Batch Utils', () => {
         publicKeyEIP7702: PUBLIC_KEY_MOCK,
       });
 
-      expect(result).toStrictEqual([CHAIN_ID_MOCK, CHAIN_ID_2_MOCK]);
+      expect(result).toStrictEqual([
+        {
+          chainId: CHAIN_ID_MOCK,
+          delegationAddress: undefined,
+          isSupported: false,
+          upgradeContractAddress: UPGRADE_CONTRACT_ADDRESS_MOCK,
+        },
+        {
+          chainId: CHAIN_ID_2_MOCK,
+          delegationAddress: CONTRACT_ADDRESS_MOCK,
+          isSupported: true,
+          upgradeContractAddress: UPGRADE_CONTRACT_ADDRESS_MOCK,
+        },
+      ]);
     });
 
-    it('excludes chain if upgraded to different contract', async () => {
-      getEIP7702SupportedChainsMock.mockReturnValueOnce([CHAIN_ID_MOCK]);
+    it('includes only specified chain IDs', async () => {
+      getEIP7702SupportedChainsMock.mockReturnValueOnce([
+        CHAIN_ID_MOCK,
+        CHAIN_ID_2_MOCK,
+      ]);
+
+      getEIP7702UpgradeContractAddressMock.mockReturnValue(
+        UPGRADE_CONTRACT_ADDRESS_MOCK,
+      );
 
       isAccountUpgradedToEIP7702Mock.mockResolvedValueOnce({
-        isSupported: false,
+        isSupported: true,
         delegationAddress: CONTRACT_ADDRESS_MOCK,
       });
 
       const result = await isAtomicBatchSupported({
         address: FROM_MOCK,
+        chainIds: [CHAIN_ID_2_MOCK, '0xabcdef'],
         getEthQuery: GET_ETH_QUERY_MOCK,
         messenger: MESSENGER_MOCK,
         publicKeyEIP7702: PUBLIC_KEY_MOCK,
       });
 
-      expect(result).toStrictEqual([]);
+      expect(result).toStrictEqual([
+        {
+          chainId: CHAIN_ID_2_MOCK,
+          delegationAddress: CONTRACT_ADDRESS_MOCK,
+          isSupported: true,
+          upgradeContractAddress: UPGRADE_CONTRACT_ADDRESS_MOCK,
+        },
+      ]);
     });
 
     it('throws if no public key', async () => {

packages/transaction-controller/src/utils/batch.ts

@@ -33,6 +33,8 @@ import type {
   PublishHook,
   TransactionBatchRequest,
   ValidateSecurityRequest,
+  IsAtomicBatchSupportedResult,
+  IsAtomicBatchSupportedResultEntry,
 } from '../types';
 import {
   TransactionEnvelopeType,
@@ -57,8 +59,9 @@ type AddTransactionBatchRequest = {
   ) => void;
 };
 
-type IsAtomicBatchSupportedRequest = {
+type IsAtomicBatchSupportedRequestInternal = {
   address: Hex;
+  chainIds?: Hex[];
   getEthQuery: (chainId: Hex) => EthQuery;
   messenger: TransactionControllerMessenger;
   publicKeyEIP7702?: Hex;
@@ -219,10 +222,11 @@ export async function addTransactionBatch(
  * @returns The chain IDs that support atomic batch transactions.
  */
 export async function isAtomicBatchSupported(
-  request: IsAtomicBatchSupportedRequest,
-): Promise<Hex[]> {
+  request: IsAtomicBatchSupportedRequestInternal,
+): Promise<IsAtomicBatchSupportedResult> {
   const {
     address,
+    chainIds,
     getEthQuery,
     messenger,
     publicKeyEIP7702: publicKey,
@@ -233,9 +237,13 @@ export async function isAtomicBatchSupported(
   }
 
   const chainIds7702 = getEIP7702SupportedChains(messenger);
-  const chainIds: Hex[] = [];
+  const results: IsAtomicBatchSupportedResultEntry[] = [];
 
   for (const chainId of chainIds7702) {
+    if (chainIds && !chainIds.includes(chainId)) {
+      continue;
+    }
+
     const ethQuery = getEthQuery(chainId);
 
     const { isSupported, delegationAddress } = await isAccountUpgradedToEIP7702(
@@ -246,14 +254,23 @@ export async function isAtomicBatchSupported(
       ethQuery,
     );
 
-    if (!delegationAddress || isSupported) {
-      chainIds.push(chainId);
-    }
+    const upgradeContractAddress = getEIP7702UpgradeContractAddress(
+      chainId,
+      messenger,
+      publicKey,
+    );
+
+    results.push({
+      chainId,
+      delegationAddress,
+      isSupported,
+      upgradeContractAddress,
+    });
   }
 
-  log('Atomic batch supported chains', chainIds);
+  log('Atomic batch supported results', results);
 
-  return chainIds;
+  return results;
 }
 
 /**