feat(multichain-account-service): add transient state support

Author: ccharly

packages/multichain-account-service/src/MultichainAccountGroup.test.ts

@@ -44,6 +44,9 @@ function setup({
   const wallet = new MultichainAccountWallet<Bip44Account<InternalAccount>>({
     providers,
     entropySource: MOCK_WALLET_1_ENTROPY_SOURCE,
+    transientState: {
+      isAlignmentInProgress: false,
+    },
   });
 
   const group = new MultichainAccountGroup({

packages/multichain-account-service/src/MultichainAccountService.ts

@@ -7,21 +7,36 @@ import type {
   Bip44Account,
   AccountProvider,
 } from '@metamask/account-api';
+import { BaseController } from '@metamask/base-controller';
 import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api';
 import { KeyringTypes } from '@metamask/keyring-controller';
 
 import type { MultichainAccountGroup } from './MultichainAccountGroup';
+import type { MultichainAccountWalletTransientState } from './MultichainAccountWallet';
 import { MultichainAccountWallet } from './MultichainAccountWallet';
 import {
   AccountProviderWrapper,
   isAccountProviderWrapper,
 } from './providers/AccountProviderWrapper';
 import { EvmAccountProvider } from './providers/EvmAccountProvider';
 import { SolAccountProvider } from './providers/SolAccountProvider';
-import type { MultichainAccountServiceMessenger } from './types';
+import type { TransientStateMetadata } from './transient-state';
+import type {
+  MultichainAccountServiceMessenger,
+  MultichainAccountServiceTransientState,
+} from './types';
 
 export const serviceName = 'MultichainAccountService';
 
+const serviceMetadata: TransientStateMetadata<MultichainAccountServiceTransientState> =
+  {
+    accountWalletsTransientState: {
+      // Transient states are never persisted.
+      persist: false,
+      anonymous: false,
+    },
+  };
+
 /**
  * The options that {@link MultichainAccountService} takes.
  */
@@ -39,7 +54,11 @@ type AccountContext<Account extends Bip44Account<KeyringAccount>> = {
 /**
  * Service to expose multichain accounts capabilities.
  */
-export class MultichainAccountService {
+export class MultichainAccountService extends BaseController<
+  typeof serviceName,
+  MultichainAccountServiceTransientState,
+  MultichainAccountServiceMessenger
+> {
   readonly #messenger: MultichainAccountServiceMessenger;
 
   readonly #providers: AccountProvider<Bip44Account<KeyringAccount>>[];
@@ -69,6 +88,15 @@ export class MultichainAccountService {
    * providers.
    */
   constructor({ messenger, providers = [] }: MultichainAccountServiceOptions) {
+    super({
+      messenger,
+      name: serviceName,
+      metadata: serviceMetadata,
+      state: {
+        accountWalletsTransientState: {},
+      },
+    });
+
     this.#messenger = messenger;
     this.#wallets = new Map();
     this.#accountIdToContext = new Map();
@@ -150,9 +178,12 @@ export class MultichainAccountService {
 
         // This will automatically "associate" all multichain accounts for that wallet
         // (based on the accounts owned by each account providers).
+        const walletId = toMultichainAccountWalletId(entropySource);
         const wallet = new MultichainAccountWallet({
           entropySource,
           providers: this.#providers,
+          transientState:
+            this.#getMultichainAccoultWalletTransientStateProxy(walletId),
         });
         this.#wallets.set(wallet.id, wallet);
 
@@ -169,6 +200,54 @@ export class MultichainAccountService {
     }
   }
 
+  #getMultichainAccoultWalletTransientStateProxy(
+    walletId: MultichainAccountWalletId,
+  ): MultichainAccountWalletTransientState {
+    // eslint-disable-next-line consistent-this, @typescript-eslint/no-this-alias
+    const service = this;
+
+    const getTransientState = <Return>(
+      state: MultichainAccountServiceTransientState,
+      then: (transientState: MultichainAccountWalletTransientState) => Return,
+    ) => {
+      return then(state.accountWalletsTransientState[walletId]);
+    };
+
+    const setTransientState = (
+      state: MultichainAccountServiceTransientState,
+      then: (transientState: MultichainAccountWalletTransientState) => void,
+    ) => {
+      then(state.accountWalletsTransientState[walletId]);
+    };
+
+    // Create initial transient state for this wallet if it does not exists yet.
+    if (!service.state.accountWalletsTransientState[walletId]) {
+      this.update((state) => {
+        state.accountWalletsTransientState[walletId] = {
+          isAlignmentInProgress: false,
+        };
+      });
+    }
+
+    return {
+      get isAlignmentInProgress(): boolean {
+        return getTransientState(
+          service.state,
+          (transientState) => transientState.isAlignmentInProgress,
+        );
+      },
+      set isAlignmentInProgress(inProgress: boolean) {
+        service.update((state) => {
+          setTransientState(
+            state,
+            (transientState) =>
+              (transientState.isAlignmentInProgress = inProgress),
+          );
+        });
+      },
+    };
+  }
+
   #handleOnAccountAdded(account: KeyringAccount): void {
     // We completely omit non-BIP-44 accounts!
     if (!isBip44Account(account)) {
@@ -177,6 +256,7 @@ export class MultichainAccountService {
 
     let sync = true;
 
+    const walletId = toMultichainAccountWalletId(account.options.entropy.id);
     let wallet = this.#wallets.get(
       toMultichainAccountWalletId(account.options.entropy.id),
     );
@@ -185,6 +265,8 @@ export class MultichainAccountService {
       wallet = new MultichainAccountWallet({
         entropySource: account.options.entropy.id,
         providers: this.#providers,
+        transientState:
+          this.#getMultichainAccoultWalletTransientStateProxy(walletId),
       });
       this.#wallets.set(wallet.id, wallet);
 

packages/multichain-account-service/src/MultichainAccountWallet.test.ts

@@ -57,6 +57,9 @@ function setup({
   const wallet = new MultichainAccountWallet<Bip44Account<InternalAccount>>({
     providers,
     entropySource,
+    transientState: {
+      isAlignmentInProgress: false,
+    },
   });
 
   return { wallet, providers };

packages/multichain-account-service/src/MultichainAccountWallet.ts

@@ -19,6 +19,17 @@ import {
 
 import { MultichainAccountGroup } from './MultichainAccountGroup';
 
+/**
+ * Transient state for a multichain account wallet.
+ *
+ * This state WILL NEVER be persisted. It will be shared with the client's UI
+ * so the UI can reflect th current state of the wallet (aligning, discovering, etc...)
+ */
+export type MultichainAccountWalletTransientState = {
+  /** True when the wallet is being aligned (creating missing accounts on each groups). */
+  isAlignmentInProgress: boolean;
+};
+
 /**
  * A multichain account wallet that holds multiple multichain accounts (one multichain account per
  * group index).
@@ -35,19 +46,22 @@ export class MultichainAccountWallet<
 
   readonly #accountGroups: Map<number, MultichainAccountGroup<Account>>;
 
-  #isAlignmentInProgress: boolean = false;
+  readonly #transientState: MultichainAccountWalletTransientState;
 
   constructor({
     providers,
     entropySource,
+    transientState,
   }: {
     providers: AccountProvider<Account>[];
     entropySource: EntropySourceId;
+    transientState: MultichainAccountWalletTransientState;
   }) {
     this.#id = toMultichainAccountWalletId(entropySource);
     this.#providers = providers;
     this.#entropySource = entropySource;
     this.#accountGroups = new Map();
+    this.#transientState = transientState;
 
     // Initial synchronization.
     this.sync();
@@ -315,23 +329,23 @@ export class MultichainAccountWallet<
    * @returns True if alignment is in progress, false otherwise.
    */
   getIsAlignmentInProgress(): boolean {
-    return this.#isAlignmentInProgress;
+    return this.#transientState.isAlignmentInProgress;
   }
 
   /**
    * Align all multichain account groups.
    */
   async alignGroups(): Promise<void> {
-    if (this.#isAlignmentInProgress) {
+    if (this.#transientState.isAlignmentInProgress) {
       return; // Prevent concurrent alignments
     }
 
-    this.#isAlignmentInProgress = true;
+    this.#transientState.isAlignmentInProgress = true;
     try {
       const groups = this.getMultichainAccountGroups();
       await Promise.all(groups.map((g) => g.align()));
     } finally {
-      this.#isAlignmentInProgress = false;
+      this.#transientState.isAlignmentInProgress = false;
     }
   }
 
@@ -341,18 +355,18 @@ export class MultichainAccountWallet<
    * @param groupIndex - The group index to align.
    */
   async alignGroup(groupIndex: number): Promise<void> {
-    if (this.#isAlignmentInProgress) {
+    if (this.#transientState.isAlignmentInProgress) {
       return; // Prevent concurrent alignments
     }
 
-    this.#isAlignmentInProgress = true;
+    this.#transientState.isAlignmentInProgress = true;
     try {
       const group = this.getMultichainAccountGroup(groupIndex);
       if (group) {
         await group.align();
       }
     } finally {
-      this.#isAlignmentInProgress = false;
+      this.#transientState.isAlignmentInProgress = false;
     }
   }
 }

packages/multichain-account-service/src/transient-state.ts

@@ -0,0 +1,16 @@
+import type {
+  StateConstraint,
+  StatePropertyMetadata,
+} from '@metamask/base-controller';
+
+/**
+ * Transient state metadata.
+ *
+ * This metadata describes how to get an anonymized representation of the state.
+ */
+export type TransientStateMetadata<T extends StateConstraint> = {
+  [P in keyof T]-?: StatePropertyMetadata<T[P]> & {
+    // Since the state is transient, it cannot be persisted.
+    persist: false;
+  };
+};

packages/multichain-account-service/src/types.ts

@@ -1,3 +1,4 @@
+import type { MultichainAccountWalletId } from '@metamask/account-api';
 import type {
   AccountsControllerAccountAddedEvent,
   AccountsControllerAccountRemovedEvent,
@@ -17,6 +18,15 @@ import type {
   MultichainAccountService,
   serviceName,
 } from './MultichainAccountService';
+import type { MultichainAccountWalletTransientState } from './MultichainAccountWallet';
+
+export type MultichainAccountServiceTransientState = {
+  accountWalletsTransientState: {
+    [
+      walletId: MultichainAccountWalletId
+    ]: MultichainAccountWalletTransientState;
+  };
+};
 
 export type MultichainAccountServiceGetMultichainAccountGroupAction = {
   type: `${typeof serviceName}:getMultichainAccountGroup`;