feat: add Websocket Service into Transaction confirmation and Token Balance Controllers

Author: Kriys94

packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts

@@ -29,6 +29,9 @@ import type {
   MultichainAssetsControllerGetStateAction,
   MultichainAssetsControllerAccountAssetListUpdatedEvent,
 } from '../MultichainAssetsController';
+import type {
+  AccountActivityServiceBalanceUpdatedEvent,
+} from '@metamask/backend-platform';
 
 const controllerName = 'MultichainBalancesController';
 
@@ -104,7 +107,8 @@ type AllowedEvents =
   | AccountsControllerAccountAddedEvent
   | AccountsControllerAccountRemovedEvent
   | AccountsControllerAccountBalancesUpdatesEvent
-  | MultichainAssetsControllerAccountAssetListUpdatedEvent;
+  | MultichainAssetsControllerAccountAssetListUpdatedEvent
+  | AccountActivityServiceBalanceUpdatedEvent;
 /**
  * Messenger type for the MultichainBalancesController.
  */
@@ -185,6 +189,19 @@ export class MultichainBalancesController extends BaseController<
         await this.#handleOnAccountAssetListUpdated(newAccountAssets);
       },
     );
+
+    // Subscribe to AccountActivityService balance updates
+    try {
+      this.messagingSystem.subscribe(
+        'AccountActivityService:balanceUpdated',
+        (balances: AccountBalancesUpdatedEventPayload) => {
+          this.#handleOnAccountBalancesUpdated(balances);
+        },
+      );
+    } catch (error) {
+      // AccountActivityService might not be available in all environments
+      console.log('AccountActivityService not available for balance updates:', error);
+    }
   }
 
   /**

packages/assets-controllers/src/TokenBalancesController.ts

@@ -41,6 +41,10 @@ import type {
   TokensControllerState,
   TokensControllerStateChangeEvent,
 } from './TokensController';
+import type {
+  AccountActivityServiceBalanceUpdatedEvent,
+} from '@metamask/backend-platform';
+import type { AccountBalancesUpdatedEventPayload } from '@metamask/keyring-api';
 
 const DEFAULT_INTERVAL = 180000;
 
@@ -104,7 +108,8 @@ export type AllowedEvents =
   | TokensControllerStateChangeEvent
   | PreferencesControllerStateChangeEvent
   | NetworkControllerStateChangeEvent
-  | KeyringControllerAccountRemovedEvent;
+  | KeyringControllerAccountRemovedEvent
+  | AccountActivityServiceBalanceUpdatedEvent;
 
 export type TokenBalancesControllerMessenger = RestrictedMessenger<
   typeof controllerName,
@@ -202,6 +207,18 @@ export class TokenBalancesController extends StaticIntervalPollingController<Tok
       'KeyringController:accountRemoved',
       (accountAddress: string) => this.#handleOnAccountRemoved(accountAddress),
     );
+
+    // Subscribe to AccountActivityService balance updates for EVM networks
+    try {
+      this.messagingSystem.subscribe(
+        'AccountActivityService:balanceUpdated',
+        (balances: AccountBalancesUpdatedEventPayload) =>
+          this.#handleAccountActivityBalanceUpdated(balances),
+      );
+    } catch (error) {
+      // AccountActivityService might not be available in all environments
+      console.log('AccountActivityService not available for EVM token balance updates:', error);
+    }
   }
 
   /**
@@ -305,6 +322,71 @@ export class TokenBalancesController extends StaticIntervalPollingController<Tok
     });
   }
 
+  /**
+   * Handles balance updates received from the AccountActivityService for EVM networks.
+   * Converts account IDs to addresses and updates token balances accordingly.
+   *
+   * @param balanceUpdate - The balance update from AccountActivityService containing new balances.
+   */
+  #handleAccountActivityBalanceUpdated(
+    balanceUpdate: AccountBalancesUpdatedEventPayload,
+  ): void {
+    // Convert account IDs to addresses first
+    const addressBalances: Record<string, Record<string, { unit: string; amount: string }>> = {};
+    
+    for (const [accountId, assetBalances] of Object.entries(balanceUpdate.balances)) {
+      let accountAddress: string;
+      
+      if (accountId.startsWith('0x')) {
+        // Already an address
+        accountAddress = accountId;
+      } else {
+        // Convert account ID to address
+        try {
+          const accounts = this.messagingSystem.call('AccountsController:listAccounts');
+          const account = accounts.find(acc => acc.id === accountId);
+          if (!account) {
+            continue;
+          }
+          accountAddress = account.address;
+        } catch (error) {
+          continue;
+        }
+      }
+      
+      // Only process valid EVM addresses
+      if (isStrictHexString(accountAddress.toLowerCase()) && isValidHexAddress(accountAddress)) {
+        addressBalances[accountAddress] = assetBalances;
+      }
+    }
+    
+    // Update state with address-based balances
+    this.update((state) => {
+      for (const [accountAddress, assetBalances] of Object.entries(addressBalances)) {
+        const address = accountAddress as Hex;
+        
+        // Initialize account if needed
+        if (!state.tokenBalances[address]) {
+          state.tokenBalances[address] = {};
+        }
+
+        // Process each EVM asset
+        for (const [assetType, balance] of Object.entries(assetBalances)) {
+          const evmERC20AssetMatch = assetType.match(/^eip155:(\d+)\/erc20:(0x[a-fA-F0-9]{40})$/i);
+          
+          if (evmERC20AssetMatch) {
+            const chainId = `0x${parseInt(evmERC20AssetMatch[1], 10).toString(16)}` as Hex;
+            const tokenAddress = evmERC20AssetMatch[2] as Hex;
+            
+            // Initialize chain and token if needed
+            ((state.tokenBalances[address] ??= {})[chainId] ??= {})[tokenAddress] = 
+              `0x${parseInt(balance.amount, 10).toString(16)}` as Hex;
+          }
+        }
+      }
+    });
+  }
+
   /**
    * Returns an array of chain ids that have tokens.
    * @param allTokens - The state for imported tokens across all chains.