chore: wip

MetaMask · Jan 22, 2025 · ba60fbf · ba60fbf
1 parent 409e083
commit ba60fbf
Showing 1 changed file with 40 additions and 2 deletions.
diff --git a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts
@@ -2,6 +2,7 @@ import type {
   AccountsControllerAccountAddedEvent,
   AccountsControllerAccountRemovedEvent,
   AccountsControllerListMultichainAccountsAction,
+  AccountsControllerAccountBalancesUpdatedEvent,
 } from '@metamask/accounts-controller';
 import {
   BaseController,
@@ -101,8 +102,8 @@ type AllowedActions =
  */
 type AllowedEvents =
   | AccountsControllerAccountAddedEvent
-  | AccountsControllerAccountRemovedEvent;
-
+  | AccountsControllerAccountRemovedEvent
+  | AccountsControllerAccountBalancesUpdatedEvent;
 /**
  * Messenger type for the MultichainBalancesController.
  */
@@ -115,6 +116,23 @@ export type MultichainBalancesControllerMessenger =
     AllowedEvents['type']
   >;
 
+/**
+ * Event emitted by AccountsController, when the balances of an account are updated.
+ */
+type AccountBalancesUpdatedEvent = {
+  type: 'AccountsController:accountBalancesUpdated';
+  payload: {
+    balances: {
+      [accountId: string]: {
+        [assetId: CaipAssetType]: {
+          amount: string;
+          unit: string;
+        };
+      };
+    };
+  };
+};
+
 /**
  * {@link MultichainBalancesController}'s metadata.
  *
@@ -176,6 +194,10 @@ export class MultichainBalancesController extends BaseController<
       'AccountsController:accountRemoved',
       (account) => this.#handleOnAccountRemoved(account),
     );
+    this.messagingSystem.subscribe(
+      'AccountsController:accountBalancesUpdated',
+      (balanceUpdate) => this.#handleOnAccountBalancesUpdated(balanceUpdate),
+    );
   }
 
   /**
@@ -313,6 +335,22 @@ export class MultichainBalancesController extends BaseController<
     // - https://github.com/MetaMask/core/blob/v213.0.0/packages/accounts-controller/src/AccountsController.ts#L1036-L1039
   }
 
+  /**
+   * Handles balance updates received from the AccountsController.
+   *
+   * @param balanceUpdate - The balance update event containing new balances.
+   */
+    #handleOnAccountBalancesUpdated(balanceUpdate: AccountBalancesUpdatedEvent): void {
+      this.update((state: Draft<MultichainBalancesControllerState>) => {
+        Object.entries(balanceUpdate.balances).forEach(([accountId, assetBalances]) => {
+          if (!state.balances[accountId]) {
+            state.balances[accountId] = {};
+          }
+          Object.assign(state.balances[accountId], assetBalances);
+        });
+      });
+    }
+
   /**
    * Handles changes when a new account has been removed.
    *