stash

MetaMask · Jan 29, 2025 · 7f3b14c · 7f3b14c
1 parent 5fe95fd
commit 7f3b14c
Show file tree

Hide file tree

Showing 3 changed files with 139 additions and 759 deletions.
diff --git a/app/scripts/lib/ComposableObservableStore.ts b/app/scripts/lib/ComposableObservableStore.ts
@@ -1,31 +1,29 @@
 import { ObservableStore } from '@metamask/obs-store';
 import {
   ActionConstraint,
+  BaseControllerInstance,
   ControllerMessenger,
   EventConstraint,
   getPersistentState,
   isBaseController,
-  isBaseControllerV1,
+  StateConstraint,
 } from '@metamask/base-controller';
 import { getKnownPropertyNames } from '@metamask/utils';
-import {
-  MemStoreControllers,
-  MemStoreControllersComposedState,
-} from '../../../shared/types/background';
 
 /**
  * An ObservableStore that can compose the state objects of its child stores and controllers
  */
-export default class ComposableObservableStore extends ObservableStore<
-  Partial<MemStoreControllersComposedState>
-> {
+export default class ComposableObservableStore<
+  Config extends Record<string, BaseControllerInstance>,
+  State extends Record<keyof Config, StateConstraint>,
+> extends ObservableStore<State> {
   /**
    * Describes which stores are being composed. The key is the name of the
    * store, and the value is either an ObservableStore, or a controller that
    * extends one of the two base controllers in the `@metamask/base-controller`
    * package.
    */
-  config: Partial<MemStoreControllers> = {};
+  config: Partial<Config> = {};
 
   controllerMessenger: ControllerMessenger<ActionConstraint, EventConstraint>;
 
@@ -46,12 +44,12 @@ export default class ComposableObservableStore extends ObservableStore<
     state = {},
     persist = false,
   }: {
-    config?: MemStoreControllers;
+    config?: Config;
     controllerMessenger: ControllerMessenger<ActionConstraint, EventConstraint>;
-    state?: Partial<MemStoreControllersComposedState>;
+    state?: Partial<State>;
     persist?: boolean;
   }) {
-    super(state);
+    super(state as State);
     this.persist = persist;
     this.controllerMessenger = controllerMessenger;
     if (config) {
@@ -67,53 +65,28 @@ export default class ComposableObservableStore extends ObservableStore<
    * with an `ObservableStore`-type `store` propeety, or a controller that extends one of the two base
    * controllers in the `@metamask/base-controller` package.
    */
-  updateStructure(config: MemStoreControllers) {
+  updateStructure<NewConfig extends Partial<Config>>(config: NewConfig) {
     this.config = config;
     this.removeAllListeners();
-    const initialState = getKnownPropertyNames(
-      config,
-    ).reduce<MemStoreControllersComposedState>(
+    const initialState = getKnownPropertyNames(config).reduce<State>(
       (composedState, controllerKey) => {
         const controller = config[controllerKey];
         if (!controller) {
-          throw new Error(`Undefined '${controllerKey}'`);
+          throw new Error(`Undefined '${String(controllerKey)}'`);
         }
 
-        if ('store' in controller && Boolean(controller.store?.subscribe)) {
-          const { store } = controller;
-          store.subscribe(
-            (state: MemStoreControllersComposedState[typeof controllerKey]) => {
-              this.#onStateChange(controllerKey, state);
-            },
-          );
-          // @ts-expect-error TODO: Widen `isBaseControllerV1` input type to `unknown`
-        } else if (isBaseControllerV1(controller)) {
-          controller.subscribe((state) => {
-            // @ts-expect-error V2 controller state excluded by type guard
-            this.#onStateChange(controllerKey, state);
-          });
-        }
-        // @ts-expect-error TODO: Widen `isBaseController{,V1}` input types to `unknown`
-        if (isBaseController(controller) || isBaseControllerV1(controller)) {
+        if (isBaseController(controller)) {
           try {
             this.controllerMessenger.subscribe<`${typeof controller.name}:stateChange`>(
               `${controller.name}:stateChange`,
               // @ts-expect-error TODO: Fix `handler` being typed as `never` by defining `Global{Actions,Events}` types and supplying them to `MetamaskController['controllerMessenger']`
-              (
-                state: MemStoreControllersComposedState[typeof controllerKey],
-              ) => {
-                let updatedState: Partial<
-                  MemStoreControllersComposedState[typeof controllerKey]
-                > = state;
+              (state: State[typeof controllerKey]) => {
+                let updatedState: Partial<State[typeof controllerKey]> = state;
                 if (this.persist && 'metadata' in controller) {
                   updatedState = getPersistentState(
-                    // @ts-expect-error No state object can be passed into this parameter because its type is wider than all V2 state objects.
-                    // TODO: Fix this parameter's type to be the widest subtype of V2 controller state types instead of their supertype/constraint.
                     state,
                     controller.metadata,
-                  ) as Partial<
-                    MemStoreControllersComposedState[typeof controllerKey]
-                  >;
+                  ) as Partial<State[typeof controllerKey]>;
                 }
                 this.#onStateChange(controllerKey, updatedState);
               },
@@ -125,17 +98,10 @@ export default class ComposableObservableStore extends ObservableStore<
           }
         }
 
-        let controllerState;
-        if ('store' in controller && 'subscribe' in controller.store) {
-          controllerState = controller.store.getState?.();
-        } else if ('state' in controller) {
-          controllerState = controller.state;
-        }
-
         composedState[controllerKey] =
-          this.persist && 'metadata' in controller && controller.metadata
-            ? getPersistentState(controllerState, controller.metadata)
-            : controllerState;
+          this.persist && controller.metadata
+            ? getPersistentState(controller.state, controller.metadata)
+            : controller.state;
         return composedState;
       },
       {} as never,
@@ -144,8 +110,8 @@ export default class ComposableObservableStore extends ObservableStore<
   }
 
   #onStateChange(
-    controllerKey: keyof MemStoreControllers,
-    newState: Partial<MemStoreControllersComposedState[typeof controllerKey]>,
+    controllerKey: keyof Config,
+    newState: Partial<State[typeof controllerKey]>,
   ) {
     const oldState = this.getState()[controllerKey];