Fix/balance change calculation gas sponsored

packages/transaction-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix balance change calculation for gas sponsored tx involving native token ([#7608](https://github.com/MetaMask/core/pull/7608))
+
 ## [62.9.1]
 
 ### Changed

packages/transaction-controller/src/TransactionController.test.ts

@@ -1023,6 +1023,10 @@ describe('TransactionController', () => {
 
     signMock = jest.fn().mockImplementation(async (transaction) => transaction);
     isEIP7702GasFeeTokensEnabledMock = jest.fn().mockResolvedValue(false);
+    getGasFeeTokensMock.mockResolvedValue({
+      gasFeeTokens: [],
+      isGasFeeSponsored: false,
+    });
     getBalanceChangesMock.mockResolvedValue({
       simulationData: SIMULATION_DATA_RESULT_MOCK,
     });
@@ -2310,6 +2314,8 @@ describe('TransactionController', () => {
 
         await controller.updateEditableParams(transactionMeta.id, {});
 
+        await flushPromises();
+
         expect(getBalanceChangesMock).toHaveBeenCalledTimes(2);
       });
 
@@ -2448,6 +2454,7 @@ describe('TransactionController', () => {
           chainId: MOCK_NETWORK.chainId,
           ethQuery: expect.any(Object),
           getSimulationConfig: expect.any(Function),
+          isGasFeeSponsored: false,
           nestedTransactions: undefined,
           txParams: {
             data: undefined,
@@ -2521,6 +2528,7 @@ describe('TransactionController', () => {
         expect(getBalanceChangesMock).toHaveBeenCalledWith(
           expect.objectContaining({
             getSimulationConfig: expect.any(Function),
+            isGasFeeSponsored: false,
           }),
         );
 
@@ -2700,6 +2708,11 @@ describe('TransactionController', () => {
         await flushPromises();
 
         expect(controller.state.transactions[0].isGasFeeSponsored).toBe(true);
+        expect(getBalanceChangesMock).toHaveBeenCalledWith(
+          expect.objectContaining({
+            isGasFeeSponsored: true,
+          }),
+        );
       });
 
       it('sets isGasFeeSponsored to false when transaction is not sponsored', async () => {
@@ -7501,6 +7514,7 @@ describe('TransactionController', () => {
         blockTime: 123,
         ethQuery: expect.any(Object),
         getSimulationConfig: expect.any(Function),
+        isGasFeeSponsored: false,
         nestedTransactions: undefined,
         txParams: {
           data: undefined,
@@ -7542,6 +7556,7 @@ describe('TransactionController', () => {
         blockTime: 123,
         ethQuery: expect.any(Object),
         getSimulationConfig: expect.any(Function),
+        isGasFeeSponsored: false,
         nestedTransactions: undefined,
         txParams: {
           data: undefined,

packages/transaction-controller/src/TransactionController.ts

@@ -4379,6 +4379,13 @@ export class TransactionController extends BaseController<
       this.#skipSimulationTransactionIds.has(transactionId);
 
     if (this.#isSimulationEnabled() && !isBalanceChangesSkipped) {
+      // Get gas fee tokens FIRST to determine if transaction is sponsored
+      // This needs to happen BEFORE getBalanceChanges so we can exclude gas costs
+      const gasFeeTokensResponse = await this.#getGasFeeTokens(transactionMeta);
+
+      gasFeeTokens = gasFeeTokensResponse?.gasFeeTokens ?? [];
+      isGasFeeSponsored = gasFeeTokensResponse?.isGasFeeSponsored ?? false;
+
       const balanceChangesResult = await this.#trace(
         { name: 'Simulate', parentContext: traceContext },
         () =>
@@ -4394,6 +4401,7 @@ export class TransactionController extends BaseController<
             },
             nestedTransactions,
             txParams,
+            isGasFeeSponsored,
           }),
       );
       simulationData = balanceChangesResult.simulationData;
@@ -4409,11 +4417,6 @@ export class TransactionController extends BaseController<
           isUpdatedAfterSecurityCheck: true,
         };
       }
-
-      const gasFeeTokensResponse = await this.#getGasFeeTokens(transactionMeta);
-
-      gasFeeTokens = gasFeeTokensResponse?.gasFeeTokens ?? [];
-      isGasFeeSponsored = gasFeeTokensResponse?.isGasFeeSponsored ?? false;
     }
 
     const latestTransactionMeta = this.#getTransaction(transactionId);

packages/transaction-controller/src/utils/balance-changes.test.ts

@@ -335,13 +335,43 @@ describe('Balance Change Utils', () => {
         });
       });
 
-      it('ignoring gas cost', async () => {
+      it('including gas cost for non-sponsored transactions', async () => {
         simulateTransactionsMock.mockResolvedValueOnce(
           createNativeBalanceResponse('0x3', '0x8', 2),
         );
 
         const result = await getBalanceChanges(REQUEST_MOCK);
 
+        // For non-sponsored transactions, withGas: true means gas is included in stateDiff
+        // previousBalance: 0x3, newBalance: 0x8 (already has gas deducted)
+        // difference: 0x8 - 0x3 = 0x5
+        expect(result).toStrictEqual({
+          simulationData: {
+            nativeBalanceChange: {
+              difference: '0x5',
+              isDecrease: false,
+              newBalance: '0x8',
+              previousBalance: '0x3',
+            },
+            tokenBalanceChanges: [],
+          },
+          gasUsed: undefined,
+        });
+      });
+
+      it('excluding gas cost for sponsored transactions', async () => {
+        simulateTransactionsMock.mockResolvedValueOnce(
+          createNativeBalanceResponse('0x3', '0xa', 0), // withGas: false means gas not deducted
+        );
+
+        const result = await getBalanceChanges({
+          ...REQUEST_MOCK,
+          isGasFeeSponsored: true,
+        });
+
+        // For sponsored transactions, withGas: false means gas is NOT included in stateDiff
+        // previousBalance: 0x3, newBalance: 0xa (value transfer only, no gas deducted)
+        // difference: 0xa - 0x3 = 0x7
         expect(result).toStrictEqual({
           simulationData: {
             nativeBalanceChange: {
@@ -354,6 +384,14 @@ describe('Balance Change Utils', () => {
           },
           gasUsed: undefined,
         });
+
+        // Verify that withGas: false was used for sponsored transaction
+        expect(simulateTransactionsMock).toHaveBeenCalledWith(
+          expect.any(String),
+          expect.objectContaining({
+            withGas: false,
+          }),
+        );
       });
     });
 
@@ -727,7 +765,7 @@ describe('Balance Change Utils', () => {
               },
             ],
             withDefaultBlockOverrides: true,
-            withGas: true,
+            withGas: true, // Token balance checks always use withGas: true
           },
         );
         expect(result).toStrictEqual({

packages/transaction-controller/src/utils/balance-changes.ts

@@ -57,6 +57,7 @@ export type GetBalanceChangesRequest = {
   getSimulationConfig: GetSimulationConfig;
   nestedTransactions?: NestedTransactionMetadata[];
   txParams: TransactionParams;
+  isGasFeeSponsored?: boolean;
 };
 
 type ParsedEvent = {
@@ -206,11 +207,10 @@ function getNativeBalanceChange(
     return undefined;
   }
 
-  return getSimulationBalanceChange(
-    previousBalance,
-    newBalance,
-    transactionResponse.gasCost,
-  );
+  // For sponsored transactions, withGas: false ensures stateDiff excludes gas costs.
+  // For non-sponsored transactions, withGas: true includes gas costs in stateDiff.
+  // Hence gas cost not needed here.
+  return getSimulationBalanceChange(previousBalance, newBalance);
 }
 
 /**
@@ -637,15 +637,15 @@ function extractLogs(
  *
  * @param previousBalance - The previous balance.
  * @param newBalance - The new balance.
- * @param offset - Optional offset to apply to the new balance.
+ * @param offset - Optional offset to apply to the new balance (as BN to maintain precision).
  * @returns The balance change data or undefined if unchanged.
  */
 function getSimulationBalanceChange(
   previousBalance: Hex,
   newBalance: Hex,
-  offset: number = 0,
+  offset: BN = new BN(0),
 ): SimulationBalanceChange | undefined {
-  const newBalanceBN = hexToBN(newBalance).add(new BN(offset));
+  const newBalanceBN = hexToBN(newBalance).add(offset);
   const previousBalanceBN = hexToBN(previousBalance);
   const differenceBN = newBalanceBN.sub(previousBalanceBN);
   const isDecrease = differenceBN.isNeg();
@@ -742,7 +742,7 @@ async function baseRequest({
     ...params,
     getSimulationConfig,
     transactions,
-    withGas: true,
+    withGas: !request.isGasFeeSponsored,
     withDefaultBlockOverrides: true,
     ...(blockTime && {
       blockOverrides: {