fix: preserve staked balances on Accounts API updates

packages/assets-controller/CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Bump `@metamask/phishing-controller` from `^17.3.0` to `^17.3.1` ([#9746](https://github.com/MetaMask/core/pull/9746))
 
+### Fixed
+
+- Preserve pooled-staking balances across Accounts API chain-slice updates (e.g. network switch / `replaceCoveredChainBalances`): exclude staking contract asset IDs from `AccountsApiDataSource` v5/v6 balance processing, and keep prior staked amounts when a merge replace omits them so Accounts API cannot reset staked ETH to missing/0 ([#9753](https://github.com/MetaMask/core/pull/9753))
+
 ## [13.1.0]
 
 ### Added

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

@@ -2493,6 +2493,47 @@ describe('AssetsController', () => {
       });
     });
 
+    it('preserves existing staked balances when replaceCoveredChainBalances omits them', async () => {
+      const stakingAssetId =
+        'eip155:1/erc20:0x4FEF9D741011476750A243aC70b9789a63dd47Df' as Caip19AssetId;
+      const initialState: Partial<AssetsControllerState> = {
+        assetsBalance: {
+          [MOCK_ACCOUNT_ID]: {
+            [MOCK_ASSET_ID]: { amount: '1' },
+            [MOCK_NATIVE_ASSET_ID]: { amount: '0.5' },
+            [stakingAssetId]: { amount: '1.5' },
+          },
+        },
+      };
+
+      await withController({ state: initialState }, async ({ controller }) => {
+        await controller.handleAssetsUpdate(
+          {
+            updateMode: 'merge',
+            replaceCoveredChainBalances: true,
+            assetsBalance: {
+              [MOCK_ACCOUNT_ID]: {
+                [MOCK_NATIVE_ASSET_ID]: { amount: '2' },
+              },
+            },
+          },
+          'AccountsApiDataSource',
+        );
+
+        expect(
+          controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[MOCK_ASSET_ID],
+        ).toBeUndefined();
+        expect(
+          controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[
+            MOCK_NATIVE_ASSET_ID
+          ],
+        ).toStrictEqual({ amount: '2' });
+        expect(
+          controller.state.assetsBalance[MOCK_ACCOUNT_ID]?.[stakingAssetId],
+        ).toStrictEqual({ amount: '1.5' });
+      });
+    });
+
     it('replaces state when full update has authoritative data', async () => {
       const initialState: Partial<AssetsControllerState> = {
         assetsBalance: {

packages/assets-controller/src/AssetsController.ts

@@ -81,6 +81,7 @@ import type {
 import { AccountActivityDataSource } from './data-sources/AccountActivityDataSource.js';
 import type { AccountsApiDataSourceConfig } from './data-sources/AccountsApiDataSource.js';
 import { AccountsApiDataSource } from './data-sources/AccountsApiDataSource.js';
+import { isStakingContractAssetId } from './data-sources/evm-rpc-services/index.js';
 import { shouldSkipNativeForCaipChainId } from './data-sources/evm-rpc-services/utils/assets.js';
 import type { PriceDataSourceConfig } from './data-sources/PriceDataSource.js';
 import {
@@ -608,6 +609,18 @@ function mergeAccountBalances(
     }
   }
 
+  // Staked vault balances are owned by StakedBalanceDataSource. When an
+  // Accounts API (or other) chain-slice replace omits them, keep the prior
+  // on-chain staked amount instead of clearing it to missing/0.
+  for (const [assetId, balance] of Object.entries(previousBalances)) {
+    if (
+      isStakingContractAssetId(assetId) &&
+      !Object.prototype.hasOwnProperty.call(next, assetId)
+    ) {
+      next[assetId] = balance;
+    }
+  }
+
   return next;
 }
 

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

@@ -565,6 +565,33 @@ describe('AccountsApiDataSource', () => {
     controller.destroy();
   });
 
+  it('excludes staking contract asset IDs from v5 balance response', async () => {
+    const stakingAssetId =
+      'eip155:1/erc20:0x4fef9d741011476750a243ac70b9789a63dd47df';
+    const balances = [
+      createMockBalanceItem(
+        `eip155:1:${MOCK_ADDRESS}`,
+        'eip155:1/slip44:60',
+        '1000000000000000000',
+      ),
+      createMockBalanceItem(`eip155:1:${MOCK_ADDRESS}`, stakingAssetId, '0'),
+    ];
+
+    const { controller } = await setupController({ balances });
+
+    const response = await controller.fetch(createDataRequest());
+    const accountBalances = response.assetsBalance?.['mock-account-id'] ?? {};
+
+    expect(accountBalances).toHaveProperty('eip155:1/slip44:60');
+    expect(
+      Object.keys(accountBalances).some((id) =>
+        id.toLowerCase().includes('0x4fef9d741011476750a243ac70b9789a63dd47df'),
+      ),
+    ).toBe(false);
+
+    controller.destroy();
+  });
+
   it('fetch marks unprocessed networks as errors', async () => {
     const { controller } = await setupController({
       unprocessedNetworks: ['eip155:1'],
@@ -755,6 +782,40 @@ describe('AccountsApiDataSource', () => {
       controller.destroy();
     });
 
+    it('excludes staking contract asset IDs from v6 balance response', async () => {
+      const stakingAssetId =
+        'eip155:1/erc20:0x4fef9d741011476750a243ac70b9789a63dd47df';
+      const { controller } = await setupController({
+        remoteFeatureFlags: { assetsAccountsApiV6: { value: true } },
+        v6Accounts: [
+          {
+            accountId: `eip155:1:${MOCK_ADDRESS}`,
+            balances: [
+              createMockV6BalanceItem(
+                'eip155:1/slip44:60',
+                '1000000000000000000',
+              ),
+              createMockV6BalanceItem(stakingAssetId, '0'),
+            ],
+          },
+        ],
+      });
+
+      const response = await controller.fetch(createDataRequest());
+      const accountBalances = response.assetsBalance?.['mock-account-id'] ?? {};
+
+      expect(accountBalances).toHaveProperty('eip155:1/slip44:60');
+      expect(
+        Object.keys(accountBalances).some((id) =>
+          id
+            .toLowerCase()
+            .includes('0x4fef9d741011476750a243ac70b9789a63dd47df'),
+        ),
+      ).toBe(false);
+
+      controller.destroy();
+    });
+
     it('marks v6 unprocessed networks as errors', async () => {
       const { controller } = await setupController({
         remoteFeatureFlags: { assetsAccountsApiV6: { value: true } },

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

@@ -34,6 +34,7 @@ import type {
   SubscriptionRequest,
 } from './AbstractDataSource.js';
 import { AbstractDataSource } from './AbstractDataSource.js';
+import { isStakingContractAssetId } from './evm-rpc-services/index.js';
 
 // ============================================================================
 // CONSTANTS
@@ -612,6 +613,13 @@ export class AccountsApiDataSource extends AbstractDataSource<
       // Normalize asset ID (checksum EVM addresses for ERC20 tokens)
       const normalizedAssetId = normalizeAssetId(item.assetId as Caip19AssetId);
 
+      // Staked balances are owned by StakedBalanceDataSource. Accounts API may
+      // return the vault share token as a normal ERC-20 (often 0 or stale),
+      // which would overwrite or wipe the on-chain staked amount on merge.
+      if (isStakingContractAssetId(normalizedAssetId)) {
+        continue;
+      }
+
       // Store balance as returned by API
       assetsBalance[accountId][normalizedAssetId] = {
         amount: item.balance,
@@ -678,6 +686,13 @@ export class AccountsApiDataSource extends AbstractDataSource<
           item.assetId as Caip19AssetId,
         );
 
+        // Staked balances are owned by StakedBalanceDataSource. Accounts API may
+        // return the vault share token as a normal ERC-20 (often 0 or stale),
+        // which would overwrite or wipe the on-chain staked amount on merge.
+        if (isStakingContractAssetId(normalizedAssetId)) {
+          continue;
+        }
+
         // Store balance as returned by API
         assetsBalance[accountId][normalizedAssetId] = {
           amount: item.balance,