[BX-944] State Restoration: Persistence (#852)

src/core/state/popupInstances/index.ts

@@ -0,0 +1,133 @@
+import { Address } from 'wagmi';
+import create from 'zustand';
+
+import { ParsedSearchAsset } from '~/core/types/assets';
+import { ChainId } from '~/core/types/chains';
+import { isNativePopup } from '~/core/utils/tabs';
+import { IndependentField } from '~/entries/popup/hooks/swap/useSwapInputs';
+import { Tab } from '~/entries/popup/pages/home';
+
+import { createStore } from '../internal/createStore';
+
+type SendAddress = Address | 'eth' | '';
+
+interface PopupInstance {
+  activeTab: Tab;
+  sendAddress: Address | string | null;
+  sendAmount: string | null;
+  sendField: 'asset' | 'native';
+  sendTokenAddressAndChain: { address: SendAddress; chainId: ChainId } | null;
+  swapAmount: string | null;
+  swapField: IndependentField | null;
+  swapTokenToBuy: ParsedSearchAsset | null;
+  swapTokenToSell: ParsedSearchAsset | null;
+}
+
+const DEFAULT_POPUP_INSTANCE_VALUES: PopupInstance = {
+  activeTab: 'tokens',
+  sendAddress: null,
+  sendAmount: null,
+  sendField: 'asset',
+  sendTokenAddressAndChain: null,
+  swapAmount: null,
+  swapField: null,
+  swapTokenToBuy: null,
+  swapTokenToSell: null,
+};
+
+export interface PopupInstanceStore extends PopupInstance {
+  resetValues: () => Promise<void>;
+  saveActiveTab: ({ tab }: { tab: Tab }) => Promise<void>;
+  saveSendAddress: ({
+    address,
+  }: {
+    address: Address | string;
+  }) => Promise<void>;
+  saveSendAmount: ({ amount }: { amount: string }) => Promise<void>;
+  saveSendField: ({ field }: { field: 'asset' | 'native' }) => Promise<void>;
+  saveSendTokenAddressAndChain: ({
+    address,
+    chainId,
+  }: {
+    address: SendAddress;
+    chainId: ChainId;
+  }) => Promise<void>;
+  saveSwapAmount: ({ amount }: { amount: string }) => Promise<void>;
+  saveSwapField: ({ field }: { field: IndependentField }) => Promise<void>;
+  saveSwapTokenToBuy: ({
+    token,
+  }: {
+    token: ParsedSearchAsset | null;
+  }) => Promise<void>;
+  saveSwapTokenToSell: ({
+    token,
+  }: {
+    token: ParsedSearchAsset | null;
+  }) => Promise<void>;
+  setupPort: () => Promise<void>;
+}
+
+export const popupInstanceStore = createStore<PopupInstanceStore>(
+  (set) => ({
+    ...DEFAULT_POPUP_INSTANCE_VALUES,
+    resetValues: popupInstanceHandlerFactory(() =>
+      set(DEFAULT_POPUP_INSTANCE_VALUES),
+    ),
+    saveActiveTab: popupInstanceHandlerFactory(({ tab }) => {
+      set({ activeTab: tab });
+    }),
+    saveSendAddress: popupInstanceHandlerFactory(({ address }) => {
+      set({ sendAddress: address });
+    }),
+    saveSendAmount: popupInstanceHandlerFactory(({ amount }) => {
+      set({ sendAmount: amount });
+    }),
+    saveSendField: popupInstanceHandlerFactory(({ field }) => {
+      set({ sendField: field });
+    }),
+    saveSendTokenAddressAndChain: popupInstanceHandlerFactory(
+      ({ address, chainId }) => {
+        set({ sendTokenAddressAndChain: { address, chainId } });
+      },
+    ),
+    saveSwapAmount: popupInstanceHandlerFactory(({ amount }) => {
+      set({ swapAmount: amount });
+    }),
+    saveSwapField: popupInstanceHandlerFactory(({ field }) => {
+      set({ swapField: field });
+    }),
+    saveSwapTokenToBuy: popupInstanceHandlerFactory(({ token }) => {
+      set({ swapTokenToBuy: token });
+    }),
+    saveSwapTokenToSell: popupInstanceHandlerFactory(({ token }) => {
+      set({ swapTokenToSell: token });
+    }),
+    setupPort: popupInstanceHandlerFactory(() => {
+      chrome.runtime.connect({ name: 'popup' });
+    }),
+  }),
+  {
+    persist: {
+      name: 'popupInstance',
+      version: 0,
+    },
+  },
+);
+
+export const usePopupInstanceStore = create(popupInstanceStore);
+
+// creates handlers that only work in popup context and passes through callback types
+function popupInstanceHandlerFactory<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  THandler extends (...args: any) => any,
+>(handler: THandler) {
+  type handlerParams = Parameters<typeof handler>;
+  return async (
+    ...args: handlerParams
+  ): Promise<void | ReturnType<THandler>> => {
+    const isPopup = await isNativePopup();
+    if (isPopup) {
+      return handler(...(<[]>args));
+    }
+  };
+}

src/core/utils/tabs.ts

@@ -24,3 +24,9 @@ export const goToNewTab = ({
     // Edge sometimes returns `Tab creation is restricted in standalone sidebar mode.
   }
 };
+
+export const isNativePopup = async () => {
+  return new Promise((resolve) => {
+    chrome.tabs?.getCurrent((tab) => resolve(!tab));
+  });
+};

src/entries/background/handlers/handleDisconnect.ts

@@ -0,0 +1,12 @@
+const POPUP_INSTANCE_DATA_EXPIRY = 180000;
+export function handleDisconnect() {
+  chrome.runtime.onConnect.addListener(function (port) {
+    if (port.name === 'popup') {
+      port.onDisconnect.addListener(async function () {
+        await chrome.storage.session.set({
+          expiry: Date.now() + POPUP_INSTANCE_DATA_EXPIRY,
+        });
+      });
+    }
+  });
+}

src/entries/background/index.ts

@@ -6,6 +6,7 @@ import { initializeSentry } from '~/core/sentry';
 import { syncStores } from '~/core/state/internal/syncStores';
 import { createWagmiClient } from '~/core/wagmi';
 
+import { handleDisconnect } from './handlers/handleDisconnect';
 import { handleInstallExtension } from './handlers/handleInstallExtension';
 import { handleKeepAlive } from './handlers/handleKeepAlive';
 import { handleProviderRequest } from './handlers/handleProviderRequest';
@@ -25,6 +26,7 @@ handleProviderRequest({ popupMessenger, inpageMessenger });
 handleTabAndWindowUpdates();
 handleSetupInpage();
 handleWallets();
+handleDisconnect();
 syncStores();
 uuid4();
 initFCM();

src/entries/popup/App.tsx

@@ -22,6 +22,7 @@ import { HWRequestListener } from './components/HWRequestListener/HWRequestListe
 import { IdleTimer } from './components/IdleTimer/IdleTimer';
 import { OnboardingKeepAlive } from './components/OnboardingKeepAlive';
 import { AuthProvider } from './hooks/useAuth';
+import { useExpiryListener } from './hooks/useExpiryListener';
 import { useIsFullScreen } from './hooks/useIsFullScreen';
 import { PlaygroundComponents } from './pages/_playgrounds';
 import { RainbowConnector } from './wagmi/RainbowConnector';
@@ -37,6 +38,7 @@ const wagmiClient = createWagmiClient({
 export function App() {
   const { currentLanguage } = useCurrentLanguageStore();
   const { deviceId } = useDeviceIdStore();
+  useExpiryListener();
 
   React.useEffect(() => {
     // Disable analytics & sentry for e2e and dev mode

src/entries/popup/hooks/send/useSendInputs.ts

@@ -2,6 +2,7 @@ import { useCallback, useMemo, useRef, useState } from 'react';
 
 import { supportedCurrencies } from '~/core/references';
 import { useCurrentCurrencyStore } from '~/core/state';
+import { usePopupInstanceStore } from '~/core/state/popupInstances';
 import { ParsedAddressAsset } from '~/core/types/assets';
 import { GasFeeLegacyParams, GasFeeParams } from '~/core/types/gas';
 import {
@@ -86,13 +87,22 @@ export const useSendInputs = ({
     }
   }, [asset, currentCurrency, independentAmount, independentField]);
 
-  const assetAmount = useMemo(
-    () =>
+  const { saveSendAmount, saveSendField } = usePopupInstanceStore();
+  const assetAmount = useMemo(() => {
+    const amount =
       independentField === 'asset'
         ? independentAmount
-        : dependentAmountDisplay?.amount,
-    [dependentAmountDisplay, independentAmount, independentField],
-  );
+        : dependentAmountDisplay?.amount;
+    saveSendField({ field: independentField });
+    saveSendAmount({ amount: independentAmount });
+    return amount;
+  }, [
+    dependentAmountDisplay,
+    independentAmount,
+    independentField,
+    saveSendAmount,
+    saveSendField,
+  ]);
 
   const setInputValue = useCallback((newValue: string) => {
     if (independentFieldRef.current) {

src/entries/popup/hooks/send/useSendState.ts

@@ -2,6 +2,7 @@ import { useMemo, useState } from 'react';
 import { Address } from 'wagmi';
 
 import { useCurrentAddressStore, useCurrentCurrencyStore } from '~/core/state';
+import { usePopupInstanceStore } from '~/core/state/popupInstances';
 import { ParsedAddressAsset } from '~/core/types/assets';
 import { ChainId } from '~/core/types/chains';
 import { isNativeAsset } from '~/core/utils/chains';
@@ -22,6 +23,7 @@ export const useSendState = ({
   rawMaxAssetBalanceAmount: string;
 }) => {
   const [toAddressOrName, setToAddressOrName] = useState<Address | string>('');
+  const { saveSendAddress } = usePopupInstanceStore();
   const { currentCurrency } = useCurrentCurrencyStore();
 
   const [, setAsset] = useState<ParsedAddressAsset>();
@@ -84,6 +86,9 @@ export const useSendState = ({
     txToAddress,
     value,
     setAsset,
-    setToAddressOrName,
+    setToAddressOrName: (address: Address | string) => {
+      setToAddressOrName(address);
+      saveSendAddress({ address });
+    },
   };
 };

src/entries/popup/hooks/swap/useSwapAssets.ts

@@ -6,6 +6,7 @@ import { selectUserAssetsListByChainId } from '~/core/resources/_selectors/asset
 import { useAssets, useUserAssets } from '~/core/resources/assets';
 import { useCurrentAddressStore, useCurrentCurrencyStore } from '~/core/state';
 import { useConnectedToHardhatStore } from '~/core/state/currentSettings/connectedToHardhat';
+import { usePopupInstanceStore } from '~/core/state/popupInstances';
 import { ParsedAsset, ParsedSearchAsset } from '~/core/types/assets';
 import { ChainId } from '~/core/types/chains';
 import { SearchAsset } from '~/core/types/search';
@@ -39,7 +40,7 @@ export const useSwapAssets = () => {
   const [assetToSell, setAssetToSellState] = useState<
     ParsedSearchAsset | SearchAsset | null
   >(null);
-  const [assetToBuy, setAssetToBuy] = useState<
+  const [assetToBuy, setAssetToBuyState] = useState<
     ParsedSearchAsset | SearchAsset | null
   >(null);
 
@@ -60,6 +61,8 @@ export const useSwapAssets = () => {
   const debouncedAssetToSellFilter = useDebounce(assetToSellFilter, 200);
   const debouncedAssetToBuyFilter = useDebounce(assetToBuyFilter, 200);
 
+  const { saveSwapTokenToBuy, saveSwapTokenToSell } = usePopupInstanceStore();
+
   const { data: userAssets = [] } = useUserAssets(
     {
       address: currentAddress,
@@ -154,6 +157,14 @@ export const useSwapAssets = () => {
     });
   }, [assetToSell, assetToSellWithPrice, userAssets]);
 
+  const setAssetToBuy = useCallback(
+    (asset: ParsedSearchAsset | null) => {
+      saveSwapTokenToBuy({ token: asset });
+      setAssetToBuyState(asset);
+    },
+    [saveSwapTokenToBuy],
+  );
+
   const setAssetToSell = useCallback(
     (asset: ParsedSearchAsset | null) => {
       if (
@@ -162,12 +173,15 @@ export const useSwapAssets = () => {
         assetToBuy?.address === asset?.address &&
         assetToBuy?.chainId === asset?.chainId
       ) {
-        setAssetToBuy(prevAssetToSell === undefined ? null : prevAssetToSell);
+        setAssetToBuyState(
+          prevAssetToSell === undefined ? null : prevAssetToSell,
+        );
       }
       setAssetToSellState(asset);
+      saveSwapTokenToSell({ token: asset });
       asset?.chainId && setOutputChainId(asset?.chainId);
     },
-    [assetToBuy, prevAssetToSell],
+    [assetToBuy, prevAssetToSell, saveSwapTokenToSell],
   );
 
   return {