From 000d1c599f0a6a94a4a1e1674047fc4ab58bb878 Mon Sep 17 00:00:00 2001
From: shanejonas <jonas.shane@gmail.com>
Date: Thu, 2 Sep 2021 12:37:14 -0700
Subject: [PATCH] Added `getPersistentState` support for V2 controllers
 `persist` metadata

---
 app/scripts/lib/ComposableObservableStore.js | 11 +++++++++--
 app/scripts/metamask-controller.js           |  1 +
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/app/scripts/lib/ComposableObservableStore.js b/app/scripts/lib/ComposableObservableStore.js
index 789a12d80f4d..1563ae5487e2 100644
--- a/app/scripts/lib/ComposableObservableStore.js
+++ b/app/scripts/lib/ComposableObservableStore.js
@@ -1,4 +1,5 @@
 import { ObservableStore } from '@metamask/obs-store';
+import { getPersistentState } from '@metamask/controllers';
 
 /**
  * @typedef {import('@metamask/controllers').ControllerMessenger} ControllerMessenger
@@ -27,9 +28,11 @@ export default class ComposableObservableStore extends ObservableStore {
    *   messenger, used for subscribing to events from BaseControllerV2-based
    *   controllers.
    * @param {Object} [options.state] - The initial store state
+   * @param {boolean} [options.persist] - Wether or not to apply the persistence for v2 controllers
    */
-  constructor({ config, controllerMessenger, state }) {
+  constructor({ config, controllerMessenger, state, persist }) {
     super(state);
+    this.persist = persist;
     this.controllerMessenger = controllerMessenger;
     if (config) {
       this.updateStructure(config);
@@ -60,7 +63,11 @@ export default class ComposableObservableStore extends ObservableStore {
         this.controllerMessenger.subscribe(
           `${store.name}:stateChange`,
           (state) => {
-            this.updateState({ [key]: state });
+            let updatedState = state;
+            if (this.persist) {
+              updatedState = getPersistentState(state, config[key].metadata);
+            }
+            this.updateState({ [key]: updatedState });
           },
         );
       }
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 804a73542fd3..17405884c2d3 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -114,6 +114,7 @@ export default class MetamaskController extends EventEmitter {
     this.store = new ComposableObservableStore({
       state: initState,
       controllerMessenger: this.controllerMessenger,
+      persist: true,
     });
 
     // external connections by origin