feat: scoped per-wallet sync for AccountTreeController

packages/account-tree-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `AccountTreeController.syncWalletWithUserStorage(entropySourceId)` and the corresponding `AccountTreeController:syncWalletWithUserStorage` messenger action, which performs a bidirectional user-storage sync for a single entropy wallet (wallet metadata + groups) without iterating every local wallet. Use this in place of `syncWithUserStorage` after operations that only affect one wallet (e.g., SRP import) ([#8929](https://github.com/MetaMask/core/pull/8929))
+
 ## [7.5.0]
 
 ### Added

packages/account-tree-controller/src/AccountTreeController-method-action-types.ts

@@ -199,6 +199,25 @@ export type AccountTreeControllerSyncWithUserStorageAtLeastOnceAction = {
   handler: AccountTreeController['syncWithUserStorageAtLeastOnce'];
 };
 
+/**
+ * Enqueues a bidirectional sync with user storage for a single entropy
+ * wallet (fire-and-forget), scoped by entropy source ID. Use this in
+ * place of `syncWithUserStorage` when only one wallet's state has changed
+ * (e.g., after an SRP import) to avoid the per-wallet fanout of fetches
+ * that a full sync triggers.
+ *
+ * IMPORTANT:
+ * No-ops if a full sync is in progress (the full sync will cover this
+ * wallet). Does NOT mark the controller as having completed its first
+ * full sync.
+ *
+ * @param entropySourceId - The entropy source ID of the wallet to sync.
+ */
+export type AccountTreeControllerSyncWalletWithUserStorageAction = {
+  type: `AccountTreeController:syncWalletWithUserStorage`;
+  handler: AccountTreeController['syncWalletWithUserStorage'];
+};
+
 /**
  * Union of all AccountTreeController action types.
  */
@@ -218,4 +237,5 @@ export type AccountTreeControllerMethodActions =
   | AccountTreeControllerSetAccountGroupHiddenAction
   | AccountTreeControllerClearStateAction
   | AccountTreeControllerSyncWithUserStorageAction
-  | AccountTreeControllerSyncWithUserStorageAtLeastOnceAction;
+  | AccountTreeControllerSyncWithUserStorageAtLeastOnceAction
+  | AccountTreeControllerSyncWalletWithUserStorageAction;

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

@@ -5138,6 +5138,30 @@ describe('AccountTreeController', () => {
     });
   });
 
+  describe('syncWalletWithUserStorage', () => {
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('calls enqueueSyncForWallet on the syncing service', () => {
+      const enqueueSyncForWalletSpy = jest
+        .spyOn(BackupAndSyncService.prototype, 'enqueueSyncForWallet')
+        .mockImplementation(() => undefined);
+
+      const { controller } = setup({
+        accounts: [MOCK_HARDWARE_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      controller.syncWalletWithUserStorage('test-entropy-id');
+
+      expect(enqueueSyncForWalletSpy).toHaveBeenCalledTimes(1);
+      expect(enqueueSyncForWalletSpy).toHaveBeenCalledWith('test-entropy-id');
+    });
+  });
+
   describe('UserStorageController:stateChange subscription', () => {
     beforeEach(() => {
       jest.clearAllMocks();

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

@@ -59,6 +59,7 @@ const MESSENGER_EXPOSED_METHODS = [
   'clearState',
   'syncWithUserStorage',
   'syncWithUserStorageAtLeastOnce',
+  'syncWalletWithUserStorage',
   'init',
   'reinit',
 ] as const;
@@ -1789,6 +1790,24 @@ export class AccountTreeController extends BaseController<
     return this.#backupAndSyncService.performFullSyncAtLeastOnce();
   }
 
+  /**
+   * Enqueues a bidirectional sync with user storage for a single entropy
+   * wallet (fire-and-forget), scoped by entropy source ID. Use this in
+   * place of `syncWithUserStorage` when only one wallet's state has changed
+   * (e.g., after an SRP import) to avoid the per-wallet fanout of fetches
+   * that a full sync triggers.
+   *
+   * IMPORTANT:
+   * No-ops if a full sync is in progress (the full sync will cover this
+   * wallet). Does NOT mark the controller as having completed its first
+   * full sync.
+   *
+   * @param entropySourceId - The entropy source ID of the wallet to sync.
+   */
+  syncWalletWithUserStorage(entropySourceId: string): void {
+    this.#backupAndSyncService.enqueueSyncForWallet(entropySourceId);
+  }
+
   /**
    * Creates an backup and sync context for sync operations.
    * Used by the backup and sync service.

packages/account-tree-controller/src/backup-and-sync/analytics/traces.test.ts

@@ -7,6 +7,7 @@ describe('BackupAndSyncAnalytics - Traces', () => {
     it('contains expected trace names', () => {
       expect(TraceName).toStrictEqual({
         AccountSyncFull: 'Multichain Account Syncing - Full',
+        AccountSyncWallet: 'Multichain Account Syncing - Wallet',
       });
     });
   });

packages/account-tree-controller/src/backup-and-sync/analytics/traces.ts

@@ -6,6 +6,7 @@ import type {
 
 export const TraceName = {
   AccountSyncFull: 'Multichain Account Syncing - Full',
+  AccountSyncWallet: 'Multichain Account Syncing - Wallet',
 } as const;
 
 /**

packages/account-tree-controller/src/backup-and-sync/service/index.test.ts

@@ -5,6 +5,10 @@ import type { AccountGroupObject } from '../../group';
 import type { AccountWalletEntropyObject } from '../../wallet';
 import { getProfileId } from '../authentication';
 import type { BackupAndSyncContext } from '../types';
+import {
+  getAllGroupsFromUserStorage,
+  getWalletFromUserStorage,
+} from '../user-storage';
 // We only need to import the functions we actually spy on
 import { getLocalEntropyWallets } from '../utils';
 
@@ -20,6 +24,14 @@ const mockGetProfileId = getProfileId as jest.MockedFunction<
 >;
 const mockGetLocalEntropyWallets =
   getLocalEntropyWallets as jest.MockedFunction<typeof getLocalEntropyWallets>;
+const mockGetWalletFromUserStorage =
+  getWalletFromUserStorage as jest.MockedFunction<
+    typeof getWalletFromUserStorage
+  >;
+const mockGetAllGroupsFromUserStorage =
+  getAllGroupsFromUserStorage as jest.MockedFunction<
+    typeof getAllGroupsFromUserStorage
+  >;
 
 describe('BackupAndSync - Service - BackupAndSyncService', () => {
   let mockContext: BackupAndSyncContext;
@@ -710,4 +722,170 @@ describe('BackupAndSync - Service - BackupAndSyncService', () => {
       expect(syncExecutionCount).toBe(2); // Still only 2
     }, 15000); // Increase timeout to 15 seconds
   });
+
+  describe('enqueueSyncForWallet', () => {
+    beforeEach(() => {
+      setupMockUserStorageControllerState(true, true);
+      jest.clearAllMocks();
+      mockGetLocalEntropyWallets.mockClear();
+    });
+
+    it('returns early when backup and sync is disabled', async () => {
+      setupMockUserStorageControllerState(false, true);
+
+      backupAndSyncService.enqueueSyncForWallet('test-entropy-id');
+
+      // Give the queue a chance to run, then assert nothing happened.
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      expect(mockGetWalletFromUserStorage).not.toHaveBeenCalled();
+      expect(mockGetLocalEntropyWallets).not.toHaveBeenCalled();
+    });
+
+    it('returns early when a full sync is in progress', async () => {
+      mockContext.controller.state.isAccountTreeSyncingInProgress = true;
+
+      backupAndSyncService.enqueueSyncForWallet('test-entropy-id');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      expect(mockGetWalletFromUserStorage).not.toHaveBeenCalled();
+      expect(mockGetLocalEntropyWallets).not.toHaveBeenCalled();
+    });
+
+    it('syncs only the matched wallet, not every local wallet', async () => {
+      mockGetLocalEntropyWallets.mockReturnValue([
+        {
+          id: 'entropy:wallet-1',
+          metadata: { entropy: { id: 'entropy-id-1' } },
+        } as unknown as AccountWalletEntropyObject,
+        {
+          id: 'entropy:wallet-2',
+          metadata: { entropy: { id: 'entropy-id-2' } },
+        } as unknown as AccountWalletEntropyObject,
+      ]);
+      mockGetWalletFromUserStorage.mockResolvedValue(null);
+      mockGetAllGroupsFromUserStorage.mockResolvedValue([]);
+
+      backupAndSyncService.enqueueSyncForWallet('entropy-id-2');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // The fetches should target only the requested wallet's entropy source.
+      expect(mockGetWalletFromUserStorage).toHaveBeenCalledTimes(1);
+      expect(mockGetWalletFromUserStorage).toHaveBeenCalledWith(
+        expect.anything(),
+        'entropy-id-2',
+      );
+      expect(mockGetAllGroupsFromUserStorage).toHaveBeenCalledTimes(1);
+      expect(mockGetAllGroupsFromUserStorage).toHaveBeenCalledWith(
+        expect.anything(),
+        'entropy-id-2',
+      );
+      expect(mockGetProfileId).toHaveBeenCalledTimes(1);
+      expect(mockGetProfileId).toHaveBeenCalledWith(
+        expect.anything(),
+        'entropy-id-2',
+      );
+    });
+
+    it('does not flip hasAccountTreeSyncingSyncedAtLeastOnce', async () => {
+      mockContext.controller.state.hasAccountTreeSyncingSyncedAtLeastOnce = false;
+      mockGetLocalEntropyWallets.mockReturnValue([
+        {
+          id: 'entropy:wallet-1',
+          metadata: { entropy: { id: 'entropy-id-1' } },
+        } as unknown as AccountWalletEntropyObject,
+      ]);
+      mockGetWalletFromUserStorage.mockResolvedValue(null);
+      mockGetAllGroupsFromUserStorage.mockResolvedValue([]);
+
+      const draft: Record<string, unknown> = {
+        hasAccountTreeSyncingSyncedAtLeastOnce: false,
+        isAccountTreeSyncingInProgress: false,
+      };
+      (mockContext.controllerStateUpdateFn as jest.Mock).mockImplementation(
+        (updater: (state: typeof draft) => void) => {
+          updater(draft);
+        },
+      );
+
+      backupAndSyncService.enqueueSyncForWallet('entropy-id-1');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // Scoped sync must NOT satisfy the first-full-sync contract.
+      expect(draft.hasAccountTreeSyncingSyncedAtLeastOnce).toBe(false);
+    });
+
+    it('toggles isAccountTreeSyncingInProgress around the sync body', async () => {
+      let flagWhileRunning: boolean | undefined;
+      const draft: Record<string, unknown> = {
+        isAccountTreeSyncingInProgress: false,
+      };
+      (mockContext.controllerStateUpdateFn as jest.Mock).mockImplementation(
+        (updater: (state: typeof draft) => void) => {
+          updater(draft);
+        },
+      );
+
+      mockGetLocalEntropyWallets.mockReturnValue([
+        {
+          id: 'entropy:wallet-1',
+          metadata: { entropy: { id: 'entropy-id-1' } },
+        } as unknown as AccountWalletEntropyObject,
+      ]);
+      // Capture the flag value at the moment the sync body runs — proves it
+      // was set true before the body started.
+      (mockContext.traceFn as jest.Mock).mockImplementation(
+        async (_: unknown, fn: () => Promise<void>) => {
+          flagWhileRunning = draft.isAccountTreeSyncingInProgress as boolean;
+          await fn();
+        },
+      );
+      mockGetWalletFromUserStorage.mockResolvedValue(null);
+      mockGetAllGroupsFromUserStorage.mockResolvedValue([]);
+
+      backupAndSyncService.enqueueSyncForWallet('entropy-id-1');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      expect(flagWhileRunning).toBe(true);
+      // And reset after the sync completes.
+      expect(draft.isAccountTreeSyncingInProgress).toBe(false);
+    });
+
+    it('does not touch isAccountTreeSyncingInProgress when the wallet is not found', async () => {
+      mockGetLocalEntropyWallets.mockReturnValue([
+        {
+          id: 'entropy:wallet-1',
+          metadata: { entropy: { id: 'entropy-id-1' } },
+        } as unknown as AccountWalletEntropyObject,
+      ]);
+
+      backupAndSyncService.enqueueSyncForWallet('unknown-entropy-id');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // No wallet found → no flag toggling, no work.
+      expect(mockContext.controllerStateUpdateFn).not.toHaveBeenCalled();
+    });
+
+    it('is a no-op when no local wallet matches the entropy source ID', async () => {
+      mockGetLocalEntropyWallets.mockReturnValue([
+        {
+          id: 'entropy:wallet-1',
+          metadata: { entropy: { id: 'entropy-id-1' } },
+        } as unknown as AccountWalletEntropyObject,
+      ]);
+
+      backupAndSyncService.enqueueSyncForWallet('unknown-entropy-id');
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      expect(mockGetWalletFromUserStorage).not.toHaveBeenCalled();
+      expect(mockGetAllGroupsFromUserStorage).not.toHaveBeenCalled();
+      expect(mockGetProfileId).not.toHaveBeenCalled();
+    });
+  });
 });