diff --git a/app/src/main/java/org/mozilla/fenix/components/StateViewModel.kt b/app/src/main/java/org/mozilla/fenix/components/StoreProvider.kt
similarity index 56%
rename from app/src/main/java/org/mozilla/fenix/components/StateViewModel.kt
rename to app/src/main/java/org/mozilla/fenix/components/StoreProvider.kt
index f2bfe15fd3..6071ebb9d3 100644
--- a/app/src/main/java/org/mozilla/fenix/components/StateViewModel.kt
+++ b/app/src/main/java/org/mozilla/fenix/components/StoreProvider.kt
@@ -9,27 +9,25 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.ViewModelProviders
 import androidx.lifecycle.get
+import mozilla.components.lib.state.Action
 import mozilla.components.lib.state.State
+import mozilla.components.lib.state.Store
 
 /**
  * Generic ViewModel to wrap a State object for state restoration
  */
 @Suppress("UNCHECKED_CAST")
-class StateViewModel<T : State>(initialState: T) : ViewModel() {
-    var state: T = initialState
-        private set(value) { field = value }
-
-    fun update(state: T) { this.state = state }
-
+class StoreProvider<S : State, A : Action, T : Store<S, A>>(val store: T) : ViewModel() {
     companion object {
-        fun <S : State> get(fragment: Fragment, initialState: S): StateViewModel<S> {
+        fun <S : State, A : Action, T : Store<S, A>> get(fragment: Fragment, initialStore: T): T {
             val factory = object : ViewModelProvider.Factory {
-                override fun <T : ViewModel?> create(modelClass: Class<T>): T {
-                    return StateViewModel(initialState) as T
+                override fun <VM : ViewModel?> create(modelClass: Class<VM>): VM {
+                    return StoreProvider(initialStore) as VM
                 }
             }
 
-            return ViewModelProviders.of(fragment, factory).get()
+            val viewModel: StoreProvider<S, A, T> = ViewModelProviders.of(fragment, factory).get()
+            return viewModel.store
         }
     }
 }
diff --git a/app/src/main/java/org/mozilla/fenix/search/SearchFragment.kt b/app/src/main/java/org/mozilla/fenix/search/SearchFragment.kt
index 14a3386155..8f4bee1d6b 100644
--- a/app/src/main/java/org/mozilla/fenix/search/SearchFragment.kt
+++ b/app/src/main/java/org/mozilla/fenix/search/SearchFragment.kt
@@ -37,7 +37,7 @@ import org.mozilla.fenix.BrowserDirection
 import org.mozilla.fenix.HomeActivity
 import org.mozilla.fenix.R
 import org.mozilla.fenix.ThemeManager
-import org.mozilla.fenix.components.StateViewModel
+import org.mozilla.fenix.components.StoreProvider
 import org.mozilla.fenix.components.metrics.Event
 import org.mozilla.fenix.ext.getSpannable
 import org.mozilla.fenix.ext.requireComponents
@@ -72,9 +72,10 @@ class SearchFragment : Fragment(), BackHandler {
         val view = inflater.inflate(R.layout.fragment_search, container, false)
         val url = session?.url ?: ""
 
-        val viewModel = StateViewModel.get(
+        searchStore = StoreProvider.get(
             this,
-            SearchState(
+            Store(
+                SearchState(
                 query = url,
                 showShortcutEnginePicker = false,
                 searchEngineSource = SearchEngineSource.Default(
@@ -82,17 +83,11 @@ class SearchFragment : Fragment(), BackHandler {
                 ),
                 showSuggestions = Settings.getInstance(requireContext()).showSearchSuggestions,
                 showVisitedSitesBookmarks = Settings.getInstance(requireContext()).shouldShowVisitedSitesBookmarks,
-                session = session
+                session = session),
+                ::searchStateReducer
             )
         )
 
-        searchStore = Store(
-            viewModel.state,
-            ::searchStateReducer
-        )
-
-        searchStore.observe(this) { viewModel.update(it) }
-
         searchInteractor = SearchInteractor(
             activity as HomeActivity,
             findNavController(),