feat(account-tree-controller): add {import,export}State actions

packages/account-tree-controller/CHANGELOG.md

@@ -9,11 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add `{import,export}State` actions ([#9663](https://github.com/MetaMask/core/pull/9663))
+  - Those methods/actions can be used to export a proper snapshot of the account-tree (including secrets or not).
+  - The payload is versionned and will auto-migrate its payload if needed on the receiving end.
+  - Currently, wallet and group IDs are not the same as the local ones, mostly because local IDs are not stable and cannot be used in a cross-client context.
 - Add `getAccountFromSelectedAccountGroup` and `setSelectedAccountGroupByAccountId` methods ([#9739](https://github.com/MetaMask/core/pull/9739))
 
 ### Changed
 
-- Bump `@metamask/keyring-api` from `^23.5.0` to `^23.7.0` ([#9676](https://github.com/MetaMask/core/pull/9676))
+- Bump `@metamask/keyring-api` from `^23.5.0` to `^24.0.0` ([#9676](https://github.com/MetaMask/core/pull/9676), [#9754](https://github.com/MetaMask/core/pull/9754))
 - Bump `@metamask/accounts-controller` from `^39.0.5` to `^39.0.6` ([#9735](https://github.com/MetaMask/core/pull/9735))
 
 ## [7.5.5]

packages/account-tree-controller/package.json

@@ -57,22 +57,24 @@
   "dependencies": {
     "@metamask/accounts-controller": "^39.0.6",
     "@metamask/base-controller": "^9.1.0",
-    "@metamask/keyring-api": "^23.7.0",
+    "@metamask/keyring-api": "^24.0.0",
     "@metamask/keyring-controller": "^27.1.0",
+    "@metamask/keyring-sdk": "^3.0.0",
     "@metamask/messenger": "^2.0.0",
     "@metamask/multichain-account-service": "^13.0.0",
     "@metamask/profile-sync-controller": "^28.3.0",
     "@metamask/snaps-controllers": "^19.0.0",
     "@metamask/snaps-sdk": "^11.0.0",
     "@metamask/snaps-utils": "^12.1.2",
-    "@metamask/superstruct": "^3.1.0",
+    "@metamask/superstruct": "^3.4.1",
     "@metamask/utils": "^11.11.0",
     "fast-deep-equal": "^3.1.3",
     "lodash": "^4.17.21"
   },
   "devDependencies": {
-    "@metamask/account-api": "^1.1.1",
+    "@metamask/account-api": "^2.0.0",
     "@metamask/auto-changelog": "^6.1.0",
+    "@metamask/eth-hd-keyring": "^15.0.0",
     "@metamask/providers": "^22.1.0",
     "@ts-bridge/cli": "^0.6.4",
     "@types/jest": "^30.0.0",

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

@@ -233,6 +233,37 @@ export type AccountTreeControllerSyncWithUserStorageAtLeastOnceAction = {
   handler: AccountTreeController['syncWithUserStorageAtLeastOnce'];
 };
 
+/**
+ * Produces a versioned snapshot of the current wallet and group state.
+ *
+ * When `options.includeSecrets` is `true` and the vault is unlocked,
+ * mnemonic phrases and private keys are included in the snapshot.
+ *
+ * @param options - Export options.
+ * @returns A promise resolving to an `AccountTreeSnapshot`.
+ */
+export type AccountTreeControllerExportStateAction = {
+  type: `AccountTreeController:exportState`;
+  handler: AccountTreeController['exportState'];
+};
+
+/**
+ * Applies a validated snapshot to the current state.
+ *
+ * Accepts an `AccountTreeSnapshot` only — untrusted wire data must be parsed
+ * with `AccountTreeSnapshot.deserialize` first. New mnemonic wallets are
+ * imported via `MultichainAccountService` and new private-key accounts via
+ * `KeyringController`. Metadata (name, pinned, hidden) is applied to all
+ * existing and newly created wallets / groups.
+ *
+ * @param snapshot - The validated snapshot to import.
+ * @returns A promise that resolves when the import is complete.
+ */
+export type AccountTreeControllerImportStateAction = {
+  type: `AccountTreeController:importState`;
+  handler: AccountTreeController['importState'];
+};
+
 /**
  * Union of all AccountTreeController action types.
  */
@@ -254,4 +285,6 @@ export type AccountTreeControllerMethodActions =
   | AccountTreeControllerSetAccountGroupHiddenAction
   | AccountTreeControllerClearStateAction
   | AccountTreeControllerSyncWithUserStorageAction
-  | AccountTreeControllerSyncWithUserStorageAtLeastOnceAction;
+  | AccountTreeControllerSyncWithUserStorageAtLeastOnceAction
+  | AccountTreeControllerExportStateAction
+  | AccountTreeControllerImportStateAction;

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

@@ -6389,4 +6389,161 @@ describe('AccountTreeController', () => {
       });
     });
   });
+
+  describe('exportState / importState round-trip', () => {
+    it('preserves wallet and group metadata across a metadata-only export/import cycle', async () => {
+      const { controller, messenger } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      const walletId = toMultichainAccountWalletId(
+        MOCK_HD_KEYRING_1.metadata.id,
+      );
+      const groupId = toMultichainAccountGroupId(
+        walletId,
+        MOCK_HD_ACCOUNT_1.options.entropy.groupIndex,
+      );
+
+      // Set custom metadata before export.
+      controller.setAccountWalletName(walletId, 'My Custom Wallet');
+      controller.setAccountGroupName(groupId, 'My Custom Account');
+      controller.setAccountGroupPinned(groupId, true);
+      controller.setAccountGroupHidden(groupId, false);
+
+      // Register handlers that export needs but the default setup() doesn't provide.
+      // withKeyringV2Unsafe: returns the entropy source ID derived from the keyring.
+      messenger.registerActionHandler(
+        'KeyringController:withKeyringV2Unsafe',
+        async (
+          _selector: unknown,
+          callback: (ctx: { keyring: unknown }) => unknown,
+        ) =>
+          callback({
+            keyring: {
+              toEntropySourceId: async () => MOCK_HD_KEYRING_1.metadata.id,
+              mnemonic: null,
+            },
+          }),
+      );
+
+      // --- EXPORT ---
+      const snapshot = await controller.exportState();
+      const payload = snapshot.serialize();
+
+      expect(payload.data.wallets).toHaveLength(1);
+      const exportedWallet = payload.data.wallets[0];
+      expect(exportedWallet.type).toBe('mnemonic');
+      expect(exportedWallet.metadata.name).toBe('My Custom Wallet');
+      expect(exportedWallet.groups[0]?.metadata.name).toBe('My Custom Account');
+      expect(exportedWallet.groups[0]?.metadata.pinned).toBe(true);
+      expect(exportedWallet.groups[0]?.metadata.hidden).toBe(false);
+
+      // The snapshot's idMap bridges local IDs ↔ payload IDs.
+      expect(snapshot.toPayloadId(walletId)).toBe(
+        `wallet:${MOCK_HD_KEYRING_1.metadata.id}`,
+      );
+      expect(
+        snapshot.toLocalId(`wallet:${MOCK_HD_KEYRING_1.metadata.id}`),
+      ).toBe(walletId);
+
+      // Mutate metadata so the import can restore it.
+      controller.setAccountWalletName(walletId, 'Overwritten Wallet Name');
+      controller.setAccountGroupName(groupId, 'Overwritten Account Name');
+      controller.setAccountGroupPinned(groupId, false);
+      controller.setAccountGroupHidden(groupId, true);
+
+      expect(
+        controller.state.accountTree.wallets[walletId]?.metadata.name,
+      ).toBe('Overwritten Wallet Name');
+
+      // --- IMPORT ---
+      // withKeyringV2Unsafe is called again during import to find the matching wallet.
+      // It's already registered; the existing handler stays in place.
+      await controller.importState(snapshot);
+
+      // After import, original metadata should be restored.
+      expect(
+        controller.state.accountTree.wallets[walletId]?.metadata.name,
+      ).toBe('My Custom Wallet');
+      expect(
+        controller.state.accountTree.wallets[walletId]?.groups[groupId]
+          ?.metadata.name,
+      ).toBe('My Custom Account');
+      expect(
+        controller.state.accountTree.wallets[walletId]?.groups[groupId]
+          ?.metadata.pinned,
+      ).toBe(true);
+      expect(
+        controller.state.accountTree.wallets[walletId]?.groups[groupId]
+          ?.metadata.hidden,
+      ).toBe(false);
+    });
+
+    it('round-trips a snapshot with includeSecrets: false and vault locked', async () => {
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      // Override KeyringController:getState to report a locked vault.
+      mocks.KeyringController.getState.mockReturnValue({
+        isUnlocked: false,
+        keyrings: mocks.KeyringController.keyrings,
+      });
+
+      messenger.registerActionHandler(
+        'KeyringController:withKeyringV2Unsafe',
+        async (
+          _selector: unknown,
+          callback: (ctx: { keyring: unknown }) => unknown,
+        ) =>
+          callback({
+            keyring: {
+              toEntropySourceId: async () => MOCK_HD_KEYRING_1.metadata.id,
+              mnemonic: null,
+            },
+          }),
+      );
+
+      // Export without secrets is allowed even when the vault is locked.
+      const snapshot = await controller.exportState({ includeSecrets: false });
+      const payload = snapshot.serialize();
+
+      // Exported mnemonic wallet has no secret value.
+      expect((payload.data.wallets[0] as { value?: string }).value).toBeUndefined();
+
+      // Reimport is a no-op for metadata when nothing changed.
+      expect(await controller.importState(snapshot)).toBeUndefined();
+    });
+
+    it('throws when exporting with includeSecrets: true and the vault is locked', async () => {
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      mocks.KeyringController.getState.mockReturnValue({
+        isUnlocked: false,
+        keyrings: mocks.KeyringController.keyrings,
+      });
+
+      messenger.registerActionHandler(
+        'KeyringController:withKeyringV2Unsafe',
+        async () => undefined,
+      );
+
+      await expect(
+        controller.exportState({ includeSecrets: true }),
+      ).rejects.toThrow(
+        'Cannot include secrets in export when vault is locked',
+      );
+    });
+  });
 });

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

@@ -36,6 +36,10 @@ import type { Rule } from './rule.js';
 import { EntropyRule } from './rules/entropy.js';
 import { KeyringRule } from './rules/keyring.js';
 import { SnapRule } from './rules/snap.js';
+import { exportState } from './state/export.js';
+import { importState } from './state/import.js';
+import type { ExportStateOptions } from './state/payload.js';
+import type { AccountTreeSnapshot } from './state/snapshot.js';
 import type {
   AccountTreeControllerConfig,
   AccountTreeControllerInternalBackupAndSyncConfig,
@@ -65,6 +69,8 @@ const MESSENGER_EXPOSED_METHODS = [
   'syncWithUserStorageAtLeastOnce',
   'init',
   'reinit',
+  'exportState',
+  'importState',
 ] as const;
 
 const accountTreeControllerMetadata: StateMetadata<AccountTreeControllerState> =
@@ -1855,6 +1861,54 @@ export class AccountTreeController extends BaseController<
     return this.#backupAndSyncService.performFullSyncAtLeastOnce();
   }
 
+  /**
+   * Produces a versioned snapshot of the current wallet and group state.
+   *
+   * When `options.includeSecrets` is `true` and the vault is unlocked,
+   * mnemonic phrases and private keys are included in the snapshot.
+   *
+   * @param options - Export options.
+   * @returns A promise resolving to an `AccountTreeSnapshot`.
+   */
+  async exportState(
+    options?: ExportStateOptions,
+  ): Promise<AccountTreeSnapshot> {
+    return exportState(
+      { getState: () => this.state, messenger: this.messenger },
+      options,
+    );
+  }
+
+  /**
+   * Applies a validated snapshot to the current state.
+   *
+   * Accepts an {@link AccountTreeSnapshot} only — untrusted wire data must be
+   * parsed with {@link AccountTreeSnapshot.deserialize} first. Callers may
+   * filter the snapshot with {@link AccountTreeSnapshot.filterWallets},
+   * {@link AccountTreeSnapshot.filterGroups}, or
+   * {@link AccountTreeSnapshot.filterAllGroups} before importing.
+   *
+   * New mnemonic wallets are imported via `MultichainAccountService` and new
+   * private-key accounts via `KeyringController`. Metadata (name, pinned,
+   * hidden) is applied to all existing and newly created wallets / groups.
+   *
+   * @param snapshot - The validated snapshot to import.
+   * @returns A promise that resolves when the import is complete.
+   */
+  async importState(snapshot: AccountTreeSnapshot): Promise<void> {
+    return importState(
+      {
+        getState: () => this.state,
+        messenger: this.messenger,
+        setWalletName: (id, name) => this.setAccountWalletName(id, name),
+        setGroupName: (id, name) => this.setAccountGroupName(id, name, true),
+        setGroupPinned: (id, pinned) => this.setAccountGroupPinned(id, pinned),
+        setGroupHidden: (id, hidden) => this.setAccountGroupHidden(id, hidden),
+      },
+      snapshot,
+    );
+  }
+
   /**
    * Creates an backup and sync context for sync operations.
    * Used by the backup and sync service.

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

@@ -40,6 +40,8 @@ export type {
   AccountTreeControllerSyncWithUserStorageAtLeastOnceAction,
   AccountTreeControllerInitAction,
   AccountTreeControllerReinitAction,
+  AccountTreeControllerExportStateAction,
+  AccountTreeControllerImportStateAction,
 } from './AccountTreeController-method-action-types.js';
 
 export type { AccountContext } from './AccountTreeController.js';
@@ -48,3 +50,29 @@ export {
   AccountTreeController,
   getDefaultAccountTreeControllerState,
 } from './AccountTreeController.js';
+
+export type {
+  AccountTreePayload,
+  AccountTreePayloadStructType,
+  AccountWalletMnemonicPayload,
+  AccountWalletPrivateKeyPayload,
+  AccountWalletMnemonicGroupEntry,
+  AccountWalletPrivateKeyGroupEntry,
+  AccountWalletPayloadId,
+  AccountGroupPayloadId,
+  AccountTreeSnapshotWallet,
+  AccountTreeSnapshotGroup,
+  ExportStateOptions,
+} from './state/payload.js';
+
+export {
+  AccountTreePayloadStruct,
+  assertValidAccountTreePayload,
+  migrate,
+  migrations,
+} from './state/payload.js';
+
+export type { VersionedState } from '@metamask/keyring-sdk';
+
+export { AccountTreeSnapshot } from './state/snapshot.js';
+export { IdMap } from './state/id-map.js';