From 7d074a0fe3cb7a3084a86b0581f194ebe0401e2f Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Tue, 18 Jun 2024 10:55:46 +0200
Subject: [PATCH] fix: Improve stability of initial state sync in Snaps UI

---
 .../snap-ui-renderer/snap-ui-renderer.js      | 20 +++++++++++++++----
 ui/contexts/snaps/snap-interface.tsx          | 16 +++++++--------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/ui/components/app/snaps/snap-ui-renderer/snap-ui-renderer.js b/ui/components/app/snaps/snap-ui-renderer/snap-ui-renderer.js
index 4b5f19e0a014..a402e6ad8c7d 100644
--- a/ui/components/app/snaps/snap-ui-renderer/snap-ui-renderer.js
+++ b/ui/components/app/snaps/snap-ui-renderer/snap-ui-renderer.js
@@ -7,7 +7,7 @@ import MetaMaskTemplateRenderer from '../../metamask-template-renderer/metamask-
 import { SnapDelineator } from '../snap-delineator';
 import {
   getSnapMetadata,
-  getMemoizedInterfaceContent,
+  getMemoizedInterface,
 } from '../../../../selectors';
 import { Box, FormTextField } from '../../../component-library';
 import { DelineatorType } from '../../../../helpers/constants/snaps';
@@ -35,8 +35,15 @@ const SnapUIRendererComponent = ({
     getSnapMetadata(state, snapId),
   );
 
-  const content = useSelector((state) =>
-    getMemoizedInterfaceContent(state, interfaceId),
+  const {
+    state: initialState,
+    context,
+    content,
+  } = useSelector(
+    (state) => getMemoizedInterface(state, interfaceId),
+    // We only want to update the state if the content has changed.
+    // We do this to avoid useless re-renders.
+    (oldState, newState) => isEqual(oldState.content, newState.content),
   );
 
   // sections are memoized to avoid useless re-renders if one of the parents element re-renders.
@@ -74,7 +81,12 @@ const SnapUIRendererComponent = ({
       boxProps={boxProps}
     >
       <Box className="snap-ui-renderer__content">
-        <SnapInterfaceContextProvider snapId={snapId} interfaceId={interfaceId}>
+        <SnapInterfaceContextProvider
+          snapId={snapId}
+          interfaceId={interfaceId}
+          initialState={initialState}
+          context={context}
+        >
           <MetaMaskTemplateRenderer sections={sections} />
         </SnapInterfaceContextProvider>
         {isPrompt && (
diff --git a/ui/contexts/snaps/snap-interface.tsx b/ui/contexts/snaps/snap-interface.tsx
index 194128af4ebe..1943e35a2745 100644
--- a/ui/contexts/snaps/snap-interface.tsx
+++ b/ui/contexts/snaps/snap-interface.tsx
@@ -3,6 +3,7 @@ import {
   InterfaceState,
   UserInputEventType,
 } from '@metamask/snaps-sdk';
+import { Json } from '@metamask/utils';
 import { debounce, throttle } from 'lodash';
 import React, {
   FunctionComponent,
@@ -11,8 +12,7 @@ import React, {
   useEffect,
   useRef,
 } from 'react';
-import { useDispatch, useSelector } from 'react-redux';
-import { getMemoizedInterface } from '../../selectors';
+import { useDispatch } from 'react-redux';
 import {
   handleSnapRequest,
   updateInterfaceState,
@@ -47,6 +47,8 @@ export const SnapInterfaceContext =
 export type SnapInterfaceContextProviderProps = {
   interfaceId: string;
   snapId: string;
+  initialState: Record<string, string | Record<string, unknown> | unknown>;
+  context: Json;
 };
 
 // We want button clicks to be instant and therefore use throttling
@@ -64,18 +66,14 @@ const THROTTLED_EVENTS = [
  * @param params.children - The childrens to wrap with the context provider.
  * @param params.interfaceId - The interface ID to use.
  * @param params.snapId - The Snap ID that requested the interface.
+ * @param params.initialState - The initial state of the interface.
+ * @param params.context - The context blob of the interface.
  * @returns The context provider.
  */
 export const SnapInterfaceContextProvider: FunctionComponent<
   SnapInterfaceContextProviderProps
-> = ({ children, interfaceId, snapId }) => {
+> = ({ children, interfaceId, snapId, initialState, context }) => {
   const dispatch = useDispatch();
-  const { state: initialState, context } = useSelector(
-    (state) => getMemoizedInterface(state, interfaceId),
-    // Prevents the selector update.
-    // We do this to avoid useless re-renders.
-    () => true,
-  );
 
   // We keep an internal copy of the state to speed-up the state update in the UI.
   // It's kept in a ref to avoid useless re-rendering of the entire tree of components.