Feature/integrate web socket service

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

@@ -29,6 +29,10 @@ import type {
   MultichainAssetsControllerGetStateAction,
   MultichainAssetsControllerAccountAssetListUpdatedEvent,
 } from '../MultichainAssetsController';
+import type {
+  AccountActivityServiceBalanceUpdatedEvent,
+  IncomingBalanceUpdate,
+} from '@metamask/backend-platform';
 
 const controllerName = 'MultichainBalancesController';
 
@@ -104,7 +108,8 @@ type AllowedEvents =
   | AccountsControllerAccountAddedEvent
   | AccountsControllerAccountRemovedEvent
   | AccountsControllerAccountBalancesUpdatesEvent
-  | MultichainAssetsControllerAccountAssetListUpdatedEvent;
+  | MultichainAssetsControllerAccountAssetListUpdatedEvent
+  | AccountActivityServiceBalanceUpdatedEvent;
 /**
  * Messenger type for the MultichainBalancesController.
  */
@@ -185,6 +190,19 @@ export class MultichainBalancesController extends BaseController<
         await this.#handleOnAccountAssetListUpdated(newAccountAssets);
       },
     );
+
+    // Subscribe to AccountActivityService balance updates
+    try {
+      this.messagingSystem.subscribe(
+        'AccountActivityService:balanceUpdated',
+        (balances: IncomingBalanceUpdate[]) => {
+          this.#handleOnAccountActivityBalancesUpdated(balances);
+        },
+      );
+    } catch (error) {
+      // AccountActivityService might not be available in all environments
+      console.log('AccountActivityService not available for balance updates:', error);
+    }
   }
 
   /**
@@ -384,6 +402,49 @@ export class MultichainBalancesController extends BaseController<
     });
   }
 
+  /**
+   * Handles balance updates received from the AccountActivityService.
+   * Converts IncomingBalanceUpdate[] format to the controller's format.
+   *
+   * @param balanceUpdates - The balance updates from AccountActivityService.
+   */
+  #handleOnAccountActivityBalancesUpdated(
+    balanceUpdates: IncomingBalanceUpdate[],
+  ): void {
+    this.update((state: Draft<MultichainBalancesControllerState>) => {
+      for (const balance of balanceUpdates) {
+        // Validate required fields
+        if (!balance.address || !balance.asset?.type || !balance.asset?.unit || balance.asset?.amount === undefined) {
+          console.warn('Skipping invalid balance update:', balance);
+          continue;
+        }
+
+        // Try to find the account ID from the address
+        let accountId: string | undefined;
+        const accounts = this.#listMultichainAccounts();
+        const account = accounts.find(acc => acc.address.toLowerCase() === balance.address.toLowerCase());
+
+        if (account) {
+          accountId = account.id;
+        } else {
+          // Skip if we can't find the account
+          continue;
+        }
+
+        // Initialize account balances if needed
+        if (!state.balances[accountId]) {
+          state.balances[accountId] = {};
+        }
+
+        // Add the balance for this asset
+        state.balances[accountId][balance.asset.type] = {
+          unit: balance.asset.unit,
+          amount: balance.asset.amount,
+        };
+      }
+    });
+  }
+
   /**
    * Handles changes when a new account has been removed.
    *

packages/assets-controllers/src/TokenBalancesController.ts

@@ -41,6 +41,10 @@ import type {
   TokensControllerState,
   TokensControllerStateChangeEvent,
 } from './TokensController';
+import type {
+  AccountActivityServiceBalanceUpdatedEvent,
+  IncomingBalanceUpdate,
+} from '@metamask/backend-platform';
 
 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,22 @@ export class TokenBalancesController extends StaticIntervalPollingController<Tok
       'KeyringController:accountRemoved',
       (accountAddress: string) => this.#handleOnAccountRemoved(accountAddress),
     );
+
+    // Subscribe to AccountActivityService balance updates for EVM networks
+    try {
+      console.log('TokenBalancesController: Attempting to subscribe to AccountActivityService:balanceUpdated...');
+      this.messagingSystem.subscribe(
+        'AccountActivityService:balanceUpdated',
+        (balances: IncomingBalanceUpdate[]) => {
+          console.log('TokenBalancesController: Received AccountActivityService:balanceUpdated event!');
+          this.#handleAccountActivityBalanceUpdated(balances);
+        },
+      );
+      console.log('TokenBalancesController: Successfully subscribed to AccountActivityService:balanceUpdated');
+    } catch (error) {
+      // AccountActivityService might not be available in all environments
+      console.log('AccountActivityService not available for EVM token balance updates:', error);
+    }
   }
 
   /**
@@ -305,6 +326,124 @@ export class TokenBalancesController extends StaticIntervalPollingController<Tok
     });
   }
 
+  /**
+   * Handles balance updates received from the AccountActivityService for EVM networks.
+   * Processes IncomingBalanceUpdate[] directly and updates token balances accordingly.
+   *
+   * @param balanceUpdates - The balance updates from AccountActivityService containing new balances.
+   */
+  #handleAccountActivityBalanceUpdated(
+    balanceUpdates: IncomingBalanceUpdate[],
+  ): void {
+    console.log('TokenBalancesController: Received balance updates:', JSON.stringify(balanceUpdates, null, 2));
+
+    const currentTokenBalances = this.state.tokenBalances;
+    const updatesWithChanges: Array<{
+      address: Hex;
+      chainId: Hex;
+      tokenAddress: Hex;
+      oldBalance: Hex;
+      newBalance: Hex;
+      hasChanged: boolean;
+    }> = [];
+
+    // Process and detect changes
+    for (const balance of balanceUpdates) {
+      // Validate required fields
+      if (!balance.address || !balance.asset?.type || !balance.asset?.unit || balance.asset?.amount === undefined) {
+        console.warn('Skipping invalid balance update:', balance);
+        continue;
+      }
+
+      const accountAddress = balance.address;
+
+      // Only process valid EVM addresses
+      if (!isStrictHexString(accountAddress.toLowerCase()) || !isValidHexAddress(accountAddress)) {
+        continue;
+      }
+
+      const address = accountAddress.toLowerCase() as Hex;
+
+      // Process EVM ERC20 assets
+      const evmERC20AssetMatch = balance.asset.type.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;
+
+        // Convert balance to hex format (matching existing behavior)
+        const newBalance = `0x${parseInt(balance.asset.amount, 10).toString(16)}` as Hex;
+
+        // Check if balance has changed
+        const currentBalance = currentTokenBalances?.[address]?.[chainId]?.[tokenAddress];
+        const hasChanged = currentBalance !== newBalance;
+
+        console.log(`TokenBalancesController: Processing balance update:`, {
+          address,
+          chainId,
+          tokenAddress,
+          rawAmount: balance.asset.amount,
+          oldBalance: currentBalance,
+          newBalance,
+          currentBalance,
+          hasChanged
+        });
+
+        // Debug the state structure
+        console.log(`TokenBalancesController: Current state structure for ${address}:`, {
+          hasAccount: !!currentTokenBalances[address],
+          accountKeys: currentTokenBalances[address] ? Object.keys(currentTokenBalances[address]) : 'none',
+          hasChain: !!(currentTokenBalances[address]?.[chainId]),
+          chainKeys: currentTokenBalances[address]?.[chainId] ? Object.keys(currentTokenBalances[address][chainId]) : 'none',
+          hasToken: !!(currentTokenBalances[address]?.[chainId]?.[tokenAddress]),
+          fullState: JSON.stringify(currentTokenBalances, null, 2)
+        });
+
+        updatesWithChanges.push({
+          address,
+          chainId,
+          tokenAddress,
+          oldBalance: currentBalance,
+          newBalance,
+          hasChanged,
+        });
+      } else {
+        console.log(`TokenBalancesController: Asset type did not match ERC20 pattern:`, balance.asset.type);
+      }
+    }
+
+    // Only update if there are actual changes
+    const changedUpdates = updatesWithChanges.filter(update => update.hasChanged);
+    if (changedUpdates.length > 0) {
+      console.log(`TokenBalancesController: Updating ${changedUpdates.length} changed balances:`, 
+        changedUpdates.map(u => ({
+          address: u.address,
+          chainId: u.chainId,
+          token: u.tokenAddress,
+          oldBalance: u.oldBalance,
+          newBalance: u.newBalance
+        }))
+      );
+
+      this.update((state) => {
+        for (const update of changedUpdates) {
+          // Ensure nested structure exists
+          if (!state.tokenBalances[update.address]) {
+            state.tokenBalances[update.address] = {};
+          }
+          if (!state.tokenBalances[update.address][update.chainId]) {
+            state.tokenBalances[update.address][update.chainId] = {};
+          }
+
+          // Update the balance
+          state.tokenBalances[update.address][update.chainId][update.tokenAddress] = update.newBalance;
+        }
+      });
+    } else {
+      console.log('TokenBalancesController: No balance changes detected, skipping state update');
+    }
+  }
+
   /**
    * Returns an array of chain ids that have tokens.
    * @param allTokens - The state for imported tokens across all chains.

packages/backend-platform/CHANGELOG.md

@@ -0,0 +1,14 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [1.0.0] - Initial Release
+
+### Added
+
+- Initial release of @metamask/backend-platform 

packages/backend-platform/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) MetaMask
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.