fix: also re-sync account with the service during init

Author: ccharly

packages/account-tree-controller/src/AccountTreeController.test.ts

@@ -46,6 +46,7 @@ import {
   getAccountTreeControllerMessenger,
   getRootMessenger,
 } from '../tests/mockMessenger';
+import { MultichainAccountService } from '@metamask/multichain-account-service';
 
 // Local mock of EMPTY_ACCOUNT to avoid circular dependency
 const EMPTY_ACCOUNT_MOCK: InternalAccount = {
@@ -332,6 +333,9 @@ function setup({
       getAccount: jest.fn(),
       getSelectedMultichainAccount: jest.fn(),
     },
+    MultichainAccountService: {
+      getMultichainAccountWallets: jest.fn(),
+    },
     UserStorageController: {
       getState: jest.fn(),
       performGetStorage: jest.fn(),
@@ -428,6 +432,17 @@ function setup({
     );
   }
 
+  // Using an empty list of wallets will make sure that no multichain accounts got
+  // removed automatically removed from the tree upon calling `init`. That's what
+  // we want for almost all tests.
+  mocks.MultichainAccountService.getMultichainAccountWallets.mockReturnValue(
+    [],
+  );
+  messenger.registerActionHandler(
+    'MultichainAccountService:getMultichainAccountWallets',
+    mocks.MultichainAccountService.getMultichainAccountWallets,
+  );
+
   const accountTreeControllerMessenger =
     getAccountTreeControllerMessenger(messenger);
   const controller = new AccountTreeController({

packages/account-tree-controller/src/AccountTreeController.ts

@@ -250,10 +250,17 @@ export class AccountTreeController extends BaseController<
       },
     );
 
+    this.messenger.subscribe(
+      'MultichainAccountService:multichainAccountGroupCreated',
+      (group) => {
+        this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group);
+      },
+    );
+
     this.messenger.subscribe(
       'MultichainAccountService:multichainAccountGroupUpdated',
       (group) => {
-        this.#handleMultichainAccountWalletGroupUpdated(group);
+        this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group);
       },
     );
 
@@ -307,6 +314,17 @@ export class AccountTreeController extends BaseController<
       this.#insert(wallets, account);
     }
 
+    // Since we're building the tree out of of the account list, we don't know if those
+    // accounts got disabled at the service level. To make sure both states are "in-sync"
+    // we check each multichain account groups and remove extra accounts if there's any!
+    for (const wallet of this.messenger.call(
+      'MultichainAccountService:getMultichainAccountWallets',
+    )) {
+      for (const group of wallet.getMultichainAccountGroups()) {
+        this.#resyncAccountsFromMultichainAccountGroups(group);
+      }
+    }
+
     // Once we have the account tree, we can apply persisted metadata (names + UI states).
     let previousSelectedAccountGroupStillExists = false;
     this.update((state) => {
@@ -1252,11 +1270,15 @@ export class AccountTreeController extends BaseController<
   }
 
   /**
-   * Handles multichain account group updates.
+   * Check if an account group object is in-sync with its multichain
+   * account group counterpart.
    *
-   * @param group - Multichain account group that got updated.
+   * If not, then the extra-accounts fromt he account group object
+   * will be removed from the tree.
+   *
+   * @param group - Multichain account group that got created or updated.
    */
-  #handleMultichainAccountWalletGroupUpdated(
+  #resyncAccountsFromMultichainAccountGroups(
     group: MultichainAccountGroup<Bip44Account<KeyringAccount>>,
   ): void {
     const multichainGroupAccounts = new Set(
@@ -1281,6 +1303,17 @@ export class AccountTreeController extends BaseController<
     this.#removeAccounts(treeGroupAccountsToRemove);
   }
 
+  /**
+   * Handles multichain account group updates.
+   *
+   * @param group - Multichain account group that got updated.
+   */
+  #handleMultichainAccountWalletGroupCreatedOrUpdated(
+    group: MultichainAccountGroup<Bip44Account<KeyringAccount>>,
+  ): void {
+    this.#resyncAccountsFromMultichainAccountGroups(group);
+  }
+
   /**
    * Gets account group object.
    *

packages/account-tree-controller/src/types.ts

@@ -40,6 +40,8 @@ import type {
   AccountTreeWalletPersistedMetadata,
 } from './wallet';
 import type {
+  MultichainAccountServiceGetMultichainAccountWalletsAction,
+  MultichainAccountServiceMultichainAccountGroupCreatedEvent,
   MultichainAccountServiceMultichainAccountGroupUpdatedEvent,
   MultichainAccountServiceWalletStatusChangeEvent,
 } from '../../multichain-account-service/src/types';
@@ -130,7 +132,8 @@ export type AllowedActions =
   | UserStorageController.UserStorageControllerPerformSetStorage
   | UserStorageController.UserStorageControllerPerformBatchSetStorage
   | AuthenticationController.AuthenticationControllerGetSessionProfile
-  | MultichainAccountServiceCreateMultichainAccountGroupAction;
+  | MultichainAccountServiceCreateMultichainAccountGroupAction
+  | MultichainAccountServiceGetMultichainAccountWalletsAction;
 
 export type AccountTreeControllerActions =
   | AccountTreeControllerGetStateAction
@@ -171,6 +174,7 @@ export type AllowedEvents =
   | AccountsControllerSelectedAccountChangeEvent
   | UserStorageController.UserStorageControllerStateChangeEvent
   | MultichainAccountServiceWalletStatusChangeEvent
+  | MultichainAccountServiceMultichainAccountGroupCreatedEvent
   | MultichainAccountServiceMultichainAccountGroupUpdatedEvent;
 
 export type AccountTreeControllerEvents =

packages/account-tree-controller/tests/mockMessenger.ts

@@ -50,6 +50,7 @@ export function getAccountTreeControllerMessenger(
       'AccountsController:selectedAccountChange',
       'UserStorageController:stateChange',
       'MultichainAccountService:walletStatusChange',
+      'MultichainAccountService:multichainAccountGroupCreated',
       'MultichainAccountService:multichainAccountGroupUpdated',
     ],
     actions: [
@@ -64,6 +65,7 @@ export function getAccountTreeControllerMessenger(
       'UserStorageController:performBatchSetStorage',
       'AuthenticationController:getSessionProfile',
       'MultichainAccountService:createMultichainAccountGroup',
+      'MultichainAccountService:getMultichainAccountWallets',
       'KeyringController:getState',
       'SnapController:get',
     ],