fix(account-tree-controller): reset mappings and selected account in clearState

packages/account-tree-controller/CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Bump `@metamask/accounts-controller` from `^39.0.7` to `^39.1.0` ([#9807](https://github.com/MetaMask/core/pull/9807))
 
+### Fixed
+
+- `clearState` now resets internal mappings and resets selected account group through `:selectedAccountGroupChange` ([#9825](https://github.com/MetaMask/core/pull/9825))
+  - Consumers are (and were already) expected to handle `''` for `:selectedAccountGroupChange` (which can happens during onboarding, and now, during wallet resets).
+
 ## [7.6.1]
 
 ### Changed

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

@@ -5515,6 +5515,49 @@ describe('AccountTreeController', () => {
         getDefaultAccountTreeControllerState(),
       );
     });
+
+    it('clears in-memory reverse-lookup Maps', () => {
+      const { controller } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      expect(controller.getAccountContext(MOCK_HD_ACCOUNT_1.id)).toBeDefined();
+      expect(
+        controller.getAccountGroupObject('entropy:mock-keyring-id-1/0'),
+      ).toBeDefined();
+
+      controller.clearState();
+
+      expect(
+        controller.getAccountContext(MOCK_HD_ACCOUNT_1.id),
+      ).toBeUndefined();
+      expect(
+        controller.getAccountGroupObject('entropy:mock-keyring-id-1/0'),
+      ).toBeUndefined();
+    });
+
+    it('publishes selectedAccountGroupChange event with empty group ID', () => {
+      const { controller, messenger } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      const previousGroupId = controller.state.selectedAccountGroup;
+      const mockListener = jest.fn();
+      messenger.subscribe(
+        'AccountTreeController:selectedAccountGroupChange',
+        mockListener,
+      );
+
+      controller.clearState();
+
+      expect(mockListener).toHaveBeenCalledWith('', previousGroupId);
+    });
   });
 
   describe('backup and sync config initialization', () => {

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

@@ -1813,13 +1813,28 @@ export class AccountTreeController extends BaseController<
   clearState(): void {
     log('Clearing state');
 
+    const previousSelectedAccountGroup = this.state.selectedAccountGroup;
+
     this.update(() => {
       return {
         ...getDefaultAccountTreeControllerState(),
       };
     });
     this.#backupAndSyncService.clearState();
 
+    // Clear in-memory reverse-lookup Maps so stale data is not accessible
+    // between this call and the next init().
+    this.#accountIdToContext.clear();
+    this.#groupIdToWalletId.clear();
+
+    // Notify subscribers that the selected group has been cleared,
+    // mirroring what #setSelectedAccountGroup does on normal transitions.
+    this.messenger.publish(
+      `${controllerName}:selectedAccountGroupChange`,
+      '',
+      previousSelectedAccountGroup,
+    );
+
     // So we know we have to call `init` again.
     this.#initialized = false;
   }

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/transaction-controller` from `^69.5.1` to `^69.5.2` ([#9823](https://github.com/MetaMask/core/pull/9823))
 
+### Fixed
+
+- Properly filter empty (`''`) selected account group event ([#9825](https://github.com/MetaMask/core/pull/9825))
+
 ## [13.1.2]
 
 ### Changed

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

@@ -3415,15 +3415,38 @@ describe('AssetsController', () => {
 
   describe('account group changes', () => {
     it('handles account group change', async () => {
-      await withController(async ({ messenger }) => {
+      await withController(async ({ controller, messenger }) => {
+        const getAssetsSpy = jest
+          .spyOn(controller, 'getAssets')
+          .mockResolvedValue([]);
+
         (messenger.publish as CallableFunction)(
           'AccountTreeController:selectedAccountGroupChange',
-          undefined,
+          'entropy:mock-keyring-id-1/0',
+          '',
         );
 
         await new Promise(process.nextTick);
 
-        expect(true).toBe(true);
+        expect(getAssetsSpy).toHaveBeenCalled();
+        getAssetsSpy.mockRestore();
+      });
+    });
+
+    it('skips asset refresh when group ID is empty (onboarding or wallet reset)', async () => {
+      await withController(async ({ controller, messenger }) => {
+        const getAssetsSpy = jest.spyOn(controller, 'getAssets');
+
+        (messenger.publish as CallableFunction)(
+          'AccountTreeController:selectedAccountGroupChange',
+          '',
+          'entropy:mock-keyring-id-1/0',
+        );
+
+        await new Promise(process.nextTick);
+
+        expect(getAssetsSpy).not.toHaveBeenCalled();
+        getAssetsSpy.mockRestore();
       });
     });
   });

packages/assets-controller/src/AssetsController.ts

@@ -1103,8 +1103,8 @@ export class AssetsController extends BaseController<
     // Subscribe to account group changes (when user switches between account groups like Account 1 -> Account 2)
     this.messenger.subscribe(
       'AccountTreeController:selectedAccountGroupChange',
-      () => {
-        this.#handleAccountGroupChanged().catch(console.error);
+      (groupId) => {
+        this.#handleAccountGroupChanged(groupId).catch(console.error);
       },
     );
 
@@ -3568,7 +3568,12 @@ export class AssetsController extends BaseController<
   // EVENT HANDLERS
   // ============================================================================
 
-  async #handleAccountGroupChanged(): Promise<void> {
+  async #handleAccountGroupChanged(groupId: string): Promise<void> {
+    // The selected account group can be empty during onboarding or wallet reset.
+    if (!groupId) {
+      return;
+    }
+
     const accounts = this.#getSelectedAccounts();
 
     log('Account group changed', {