feat(account-tree-controller)!: add strip{Metadata,Secrets} helpers on state snapshot

packages/account-tree-controller/CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `AccountTreeSnapshot:strip{Metadata,Secrets}` support ([#10112](https://github.com/MetaMask/core/pull/10112))
+  - This can be used to mutate the snapshot before calling `:importState` and ease custom flow integration where secrets and metadata would be imported in 2-steps.
+
+### Changed
+
+- **BREAKING:** `metadata` is now optional on `AccountWalletPayloadMetadata`, `AccountWalletGroupPayloadMetadata`, and all wallet/group payload entry types ([#10112](https://github.com/MetaMask/core/pull/10112))
+  - Consumers reading `metadata` from a snapshot or payload must now guard against `undefined` (e.g. after calling `AccountTreeSnapshot.stripMetadata()`).
+
 ## [8.1.0]
 
 ### Added

packages/account-tree-controller/src/state/import.ts

@@ -168,9 +168,11 @@ function setGroupMetadata(
   localGroupId: AccountGroupId,
   payloadGroupMetadata: AccountWalletMnemonicGroupEntry['metadata'],
 ): void {
-  context.setAccountGroupName(localGroupId, payloadGroupMetadata.name);
-  context.setAccountGroupPinned(localGroupId, payloadGroupMetadata.pinned);
-  context.setAccountGroupHidden(localGroupId, payloadGroupMetadata.hidden);
+  if (payloadGroupMetadata) {
+    context.setAccountGroupName(localGroupId, payloadGroupMetadata.name);
+    context.setAccountGroupPinned(localGroupId, payloadGroupMetadata.pinned);
+    context.setAccountGroupHidden(localGroupId, payloadGroupMetadata.hidden);
+  }
 }
 
 /**
@@ -254,7 +256,9 @@ async function importMnemonicWallet(
     localWallet = findLocalWalletMnemonicFromId(context, id);
   }
 
-  context.setWalletName(localWallet.id, payloadWallet.metadata.name);
+  if (payloadWallet.metadata) {
+    context.setWalletName(localWallet.id, payloadWallet.metadata.name);
+  }
 
   // Apply metadata to groups that are already present locally before attempting
   // to create missing ones. If createMultichainAccountGroups throws partway

packages/account-tree-controller/src/state/payload.ts

@@ -89,7 +89,7 @@ export type AccountWalletMnemonicGroupEntry = {
   id: AccountGroupPayloadId;
   /** BIP-44 account index this group was derived at. */
   groupIndex: number;
-  metadata: AccountWalletGroupPayloadMetadata;
+  metadata?: AccountWalletGroupPayloadMetadata;
 };
 
 /**
@@ -123,7 +123,7 @@ export type AccountWalletPrivateKeyGroupEntry = {
      */
     type?: KeyringAccount['type'];
   };
-  metadata: AccountWalletGroupPayloadMetadata;
+  metadata?: AccountWalletGroupPayloadMetadata;
 };
 
 /** Payload entry for an HD (entropy) wallet and its derived account groups. */
@@ -132,7 +132,7 @@ export type AccountWalletMnemonicPayload = {
   type: typeof AccountWalletPayloadType.Mnemonic;
   /** BIP-39 mnemonic phrase encoded as bytes. Absent in metadata-only exports. */
   value?: EncodedBytes;
-  metadata: AccountWalletPayloadMetadata;
+  metadata?: AccountWalletPayloadMetadata;
   groups: AccountWalletMnemonicGroupEntry[];
 };
 
@@ -145,7 +145,7 @@ export type AccountWalletMnemonicPayload = {
 export type AccountWalletPrivateKeyPayload = {
   id: AccountWalletPayloadId;
   type: typeof AccountWalletPayloadType.PrivateKey;
-  metadata: AccountWalletPayloadMetadata;
+  metadata?: AccountWalletPayloadMetadata;
   groups: AccountWalletPrivateKeyGroupEntry[];
 };
 
@@ -257,13 +257,13 @@ const AccountWalletPrivateKeyValueStruct = object({
 const AccountWalletMnemonicGroupEntryStruct = object({
   id: AccountGroupPayloadIdStruct,
   groupIndex: integer(),
-  metadata: AccountWalletGroupPayloadMetadataStruct,
+  metadata: exactOptional(AccountWalletGroupPayloadMetadataStruct),
 });
 
 const AccountWalletPrivateKeyGroupEntryStruct = object({
   id: AccountGroupPayloadIdStruct,
   value: exactOptional(AccountWalletPrivateKeyValueStruct),
-  metadata: AccountWalletGroupPayloadMetadataStruct,
+  metadata: exactOptional(AccountWalletGroupPayloadMetadataStruct),
 });
 
 // The `groups` array in a mnemonic wallet payload must have contiguous group indices starting at 0.
@@ -289,14 +289,14 @@ const AccountWalletMnemonicPayloadStruct = object({
   id: AccountWalletPayloadIdStruct,
   type: literal(AccountWalletPayloadType.Mnemonic),
   value: exactOptional(sensitive(BytesStruct)),
-  metadata: AccountWalletPayloadMetadataStruct,
+  metadata: exactOptional(AccountWalletPayloadMetadataStruct),
   groups: AccountWalletMnemonicGroupsStruct,
 });
 
 const AccountWalletPrivateKeyPayloadStruct = object({
   id: AccountWalletPayloadIdStruct,
   type: literal(AccountWalletPayloadType.PrivateKey),
-  metadata: AccountWalletPayloadMetadataStruct,
+  metadata: exactOptional(AccountWalletPayloadMetadataStruct),
   groups: array(AccountWalletPrivateKeyGroupEntryStruct),
 });
 

packages/account-tree-controller/src/state/snapshot.test.ts

@@ -6,6 +6,7 @@ import type {
 import {
   ACCOUNT_TREE_PAYLOAD_CURRENT_VERSION,
   AccountWalletPayloadType,
+  AccountWalletPrivateKeyEncoding,
   toGroupPayloadId,
   toWalletPayloadId,
 } from './payload.js';
@@ -19,6 +20,7 @@ const MOCK_PRIVATE_KEY_PAYLOAD_ID = toWalletPayloadId(
 const MOCK_MNEMONIC_WALLET: AccountWalletMnemonicPayload = {
   id: MOCK_MNEMONIC_PAYLOAD_ID,
   type: AccountWalletPayloadType.Mnemonic,
+  value: [1, 2, 3, 4],
   metadata: { name: 'Wallet 1' },
   groups: [
     {
@@ -41,6 +43,10 @@ const MOCK_PRIVATE_KEY_WALLET: AccountWalletPrivateKeyPayload = {
   groups: [
     {
       id: toGroupPayloadId(MOCK_PRIVATE_KEY_PAYLOAD_ID, '0xdeadbeef'),
+      value: {
+        privateKey: [0xde, 0xad, 0xbe, 0xef],
+        encoding: AccountWalletPrivateKeyEncoding.Hexadecimal,
+      },
       metadata: { name: 'Imported 1', pinned: false, hidden: true },
     },
   ],
@@ -412,4 +418,64 @@ describe('AccountTreeSnapshot', () => {
       ).rejects.toThrow('Invalid AccountTreePayload');
     });
   });
+
+  describe('stripSecrets', () => {
+    it('removes value from mnemonic wallets', () => {
+      const snapshot = new AccountTreeSnapshot([MOCK_MNEMONIC_WALLET]);
+      const stripped = snapshot.stripSecrets();
+      expect(stripped.serialize().wallets[0]).not.toHaveProperty('value');
+    });
+
+    it('removes value from private-key group entries', () => {
+      const snapshot = new AccountTreeSnapshot([MOCK_PRIVATE_KEY_WALLET]);
+      const stripped = snapshot.stripSecrets();
+      expect(stripped.serialize().wallets[0]?.groups[0]).not.toHaveProperty(
+        'value',
+      );
+    });
+
+    it('preserves wallet and group metadata', () => {
+      const snapshot = new AccountTreeSnapshot([
+        MOCK_MNEMONIC_WALLET,
+        MOCK_PRIVATE_KEY_WALLET,
+      ]);
+      const stripped = snapshot.stripSecrets();
+      const { wallets } = stripped.serialize();
+      expect(wallets[0]?.metadata.name).toBe('Wallet 1');
+      expect(wallets[0]?.groups[0]?.metadata.name).toBe('Account 1');
+      expect(wallets[1]?.metadata.name).toBe('Imported Accounts');
+      expect(wallets[1]?.groups[0]?.metadata.name).toBe('Imported 1');
+    });
+  });
+
+  describe('stripMetadata', () => {
+    it('removes wallet metadata', () => {
+      const snapshot = new AccountTreeSnapshot([MOCK_MNEMONIC_WALLET]);
+      const stripped = snapshot.stripMetadata();
+      expect(stripped.serialize().wallets[0]).not.toHaveProperty('metadata');
+    });
+
+    it('removes group metadata', () => {
+      const snapshot = new AccountTreeSnapshot([MOCK_MNEMONIC_WALLET]);
+      const stripped = snapshot.stripMetadata();
+      expect(stripped.serialize().wallets[0]?.groups[0]).not.toHaveProperty(
+        'metadata',
+      );
+    });
+
+    it('preserves secret values', () => {
+      const snapshot = new AccountTreeSnapshot([
+        MOCK_MNEMONIC_WALLET,
+        MOCK_PRIVATE_KEY_WALLET,
+      ]);
+      const stripped = snapshot.stripMetadata();
+      const { wallets } = stripped.serialize();
+      expect((wallets[0] as typeof MOCK_MNEMONIC_WALLET).value).toStrictEqual(
+        MOCK_MNEMONIC_WALLET.value,
+      );
+      expect(
+        (wallets[1] as typeof MOCK_PRIVATE_KEY_WALLET).groups[0]?.value,
+      ).toStrictEqual(MOCK_PRIVATE_KEY_WALLET.groups[0]?.value);
+    });
+  });
 });

packages/account-tree-controller/src/state/snapshot.ts

@@ -72,7 +72,7 @@ export class AccountTreeSnapshot {
    * Filters groups within one wallet. Other wallets are left unchanged.
    *
    * Throws if `walletId` does not identify a wallet in the snapshot.
-   * Removes the wallet if no groups remain after filtering — this prevents a
+   * Removes the wallet if no groups remain after filtering - this prevents a
    * mnemonic wallet with zero selected groups from still transferring its secret.
    *
    * **Mnemonic wallets:** group indices must remain contiguous starting at 0
@@ -171,6 +171,58 @@ export class AccountTreeSnapshot {
     return new AccountTreeSnapshot(filteredEntries, this.#idMap);
   }
 
+  /**
+   * Returns a new snapshot with all secret material removed - mnemonic
+   * {@link AccountWalletMnemonicPayload.value | values} and private-key group
+   * {@link AccountWalletPrivateKeyGroupEntry.value | values} are omitted.
+   * Wallet and group metadata (names, pin, hidden) are preserved.
+   *
+   * Useful when only metadata (names, layout) needs to be applied and secrets
+   * are already present in the vault.
+   *
+   * @returns A secrets-stripped snapshot.
+   */
+  stripSecrets(): AccountTreeSnapshot {
+    const entries = this.#entries.map((wallet): AccountTreeWalletEntry => {
+      if (wallet.type === AccountWalletPayloadType.Mnemonic) {
+        const { value: _value, ...rest } = wallet;
+        return rest;
+      }
+      return {
+        ...wallet,
+        groups: wallet.groups.map(
+          ({ value: _value, ...group }): AccountWalletPrivateKeyGroupEntry =>
+            group,
+        ),
+      };
+    });
+    return new AccountTreeSnapshot(entries, this.#idMap);
+  }
+
+  /**
+   * Returns a new snapshot with all metadata reset to defaults - wallet names
+   * are cleared and group metadata (`name`, `pinned`, `hidden`) is reset.
+   * Secret values are preserved.
+   *
+   * Useful when importing secrets into a vault in a separate step from applying
+   * metadata - the metadata can be re-applied later from the original snapshot.
+   *
+   * @returns A metadata-stripped snapshot.
+   */
+  stripMetadata(): AccountTreeSnapshot {
+    const entries = this.#entries.map((wallet): AccountTreeWalletEntry => {
+      const { metadata: _walletMetadata, ...walletRest } = wallet;
+      return {
+        ...walletRest,
+        groups: wallet.groups.map((group) => {
+          const { metadata: _groupMetadata, ...groupRest } = group;
+          return groupRest as typeof group;
+        }),
+      } as AccountTreeWalletEntry; // Looks like the compiler is not able to infer this correctly, but we just remove the `metadata` field out of any entry.
+    });
+    return new AccountTreeSnapshot(entries, this.#idMap);
+  }
+
   /**
    * Converts a payload ID (wallet or group) to the corresponding local
    * `AccountTreeController` ID.
@@ -226,7 +278,7 @@ export class AccountTreeSnapshot {
    * versions and wallet types fail closed with an error instead of returning a
    * partial snapshot.
    *
-   * The returned snapshot has no ID map — {@link toLocalId} / {@link toPayloadId}
+   * The returned snapshot has no ID map - {@link toLocalId} / {@link toPayloadId}
    * return `undefined`. Pass an {@link IdMap} to the constructor when you need
    * the map.
    *