feat: make form state persistent between route changes (#146)

* feat(packages): make form state persistent between route changes

* refactor(packages): rename variable

* fix(packages): lint

* chore(packages): simplify

* fix(packages): lint

Author: necipsagiro

services/simple-staking/src/ui/baby/layout.tsx

@@ -112,7 +112,7 @@ function BabyLayoutContent() {
       as="main"
       className="mx-auto flex max-w-[760px] flex-1 flex-col gap-[3rem] pb-24"
     >
-      <Tabs items={fallbackTabItems} defaultActiveTab="stake" />
+      <Tabs items={fallbackTabItems} defaultActiveTab="stake" keepMounted />
     </Container>
   );
 

services/simple-staking/src/ui/baby/widgets/StakingForm/index.tsx

@@ -1,23 +1,32 @@
 import { Form } from "@babylonlabs-io/core-ui";
+import { useMemo } from "react";
+import { DeepPartial } from "react-hook-form";
 
 import { AmountField } from "@/ui/baby/components/AmountField";
 import { FeeField } from "@/ui/baby/components/FeeField";
-import { useStakingState } from "@/ui/baby/state/StakingState";
+import { useStakingState, type FormData } from "@/ui/baby/state/StakingState";
 import { StakingModal } from "@/ui/baby/widgets/StakingModal";
 import { SubmitButton } from "@/ui/baby/widgets/SubmitButton";
 import { ValidatorField } from "@/ui/baby/widgets/ValidatorField";
 import { FormAlert } from "@/ui/common/components/Multistaking/MultistakingForm/FormAlert";
-
-interface FormFields {
-  amount: number;
-  validatorAddresses: string[];
-  feeAmount: number;
-}
+import { useFormPersistenceState } from "@/ui/common/state/FormPersistenceState";
 
 interface StakingFormProps {
   isGeoBlocked?: boolean;
 }
 
+/**
+ * StakingForm supports multi-validator selection and draft persistence, so it
+ * uses 'validatorAddresses' (string[]) instead of the single 'validatorAddress'
+ * expected by 'FormData' in 'StakingState'.
+ *
+ * This interface removes 'validatorAddress' from 'FormData' and adds
+ * 'validatorAddresses' to align with the validation schema and 'FormPersistenceState'.
+ */
+export interface StakingFormFields extends Omit<FormData, "validatorAddress"> {
+  validatorAddresses: string[];
+}
+
 export default function StakingForm({
   isGeoBlocked = false,
 }: StakingFormProps) {
@@ -31,19 +40,44 @@ export default function StakingForm({
     disabled,
   } = useStakingState();
 
+  const { babyStakeDraft, setBabyStakeDraft } = useFormPersistenceState();
+
+  const defaultValues = useMemo<Partial<StakingFormFields>>(() => {
+    return {
+      amount: babyStakeDraft?.amount,
+      validatorAddresses: babyStakeDraft?.validatorAddresses,
+      feeAmount: babyStakeDraft?.feeAmount,
+    };
+  }, [babyStakeDraft]);
+
   const handlePreview = ({
     amount,
     validatorAddresses,
     feeAmount,
-  }: FormFields) => {
-    showPreview({ amount, feeAmount, validatorAddress: validatorAddresses[0] });
+  }: Required<StakingFormFields>) => {
+    showPreview({
+      amount,
+      feeAmount,
+      validatorAddress: validatorAddresses[0],
+    });
+  };
+
+  const handleChange = (data: DeepPartial<StakingFormFields>) => {
+    setBabyStakeDraft({
+      ...data,
+      validatorAddresses: data.validatorAddresses?.filter(
+        (i) => i !== undefined,
+      ),
+    });
   };
 
   return (
     <Form
       schema={formSchema}
       className="flex h-[500px] flex-col gap-2"
       onSubmit={handlePreview}
+      defaultValues={defaultValues}
+      onChange={handleChange}
     >
       <AmountField balance={availableBalance} price={babyPrice} />
       <ValidatorField />

services/simple-staking/src/ui/common/components/Multistaking/MultistakingForm/MultistakingForm.tsx

@@ -1,6 +1,8 @@
 import { Form } from "@babylonlabs-io/core-ui";
-import { useCallback } from "react";
+import { useCallback, useMemo } from "react";
+import { DeepPartial } from "react-hook-form";
 
+import { useFormPersistenceState } from "@/ui/common/state/FormPersistenceState";
 import {
   useMultistakingState,
   type MultistakingFormFields,
@@ -12,31 +14,64 @@ import { MultistakingFormContent } from "./MultistakingFormContent";
 export function MultistakingForm() {
   const { stakingInfo, setFormData, goToStep } = useStakingState();
   const { validationSchema } = useMultistakingState();
+  const { btcStakeDraft, setBtcStakeDraft } = useFormPersistenceState();
+
+  const defaultValues = useMemo<Partial<MultistakingFormFields>>(
+    () => ({
+      finalityProviders: btcStakeDraft?.finalityProviders,
+      amount: btcStakeDraft?.amount,
+      term: btcStakeDraft?.term ?? stakingInfo?.defaultStakingTimeBlocks,
+      feeRate: btcStakeDraft?.feeRate ?? stakingInfo?.defaultFeeRate ?? 0,
+      feeAmount: btcStakeDraft?.feeAmount,
+    }),
+    [
+      btcStakeDraft,
+      stakingInfo?.defaultStakingTimeBlocks,
+      stakingInfo?.defaultFeeRate,
+    ],
+  );
 
   const handlePreview = useCallback(
-    (formValues: MultistakingFormFields) => {
+    (formValues: Required<MultistakingFormFields>) => {
       setFormData({
         finalityProviders: Object.values(formValues.finalityProviders),
-        term: Number(formValues.term),
-        amount: Number(formValues.amount),
-        feeRate: Number(formValues.feeRate),
-        feeAmount: Number(formValues.feeAmount),
+        term: formValues.term,
+        amount: formValues.amount,
+        feeRate: formValues.feeRate,
+        feeAmount: formValues.feeAmount,
       });
 
       goToStep(StakingStep.PREVIEW);
     },
     [setFormData, goToStep],
   );
 
+  const handleChange = (data: DeepPartial<MultistakingFormFields>) => {
+    const sanitizedFinalityProviders: Record<string, string> = {};
+
+    Object.entries(data.finalityProviders ?? {}).forEach(([key, value]) => {
+      if (typeof value === "string") {
+        sanitizedFinalityProviders[key] = value;
+      }
+    });
+
+    setBtcStakeDraft({
+      ...data,
+      finalityProviders: sanitizedFinalityProviders,
+    });
+  };
+
   if (!stakingInfo) {
     return null;
   }
 
   return (
     <Form
-      schema={validationSchema as any}
+      schema={validationSchema}
       mode="onChange"
       reValidateMode="onChange"
+      defaultValues={defaultValues}
+      onChange={handleChange}
       onSubmit={handlePreview}
     >
       <MultistakingFormContent />

services/simple-staking/src/ui/common/components/Multistaking/MultistakingForm/StakingFeesSection.tsx

@@ -11,7 +11,6 @@ import { getNetworkConfigBTC } from "@/ui/common/config/network/btc";
 import { BBN_FEE_AMOUNT } from "@/ui/common/constants";
 import { usePrice } from "@/ui/common/hooks/client/api/usePrices";
 import { useStakingService } from "@/ui/common/hooks/services/useStakingService";
-import { useStakingState } from "@/ui/common/state/StakingState";
 import { satoshiToBtc } from "@/ui/common/utils/btc";
 import { calculateTokenValueInCurrency } from "@/ui/common/utils/formatCurrency";
 import { maxDecimals } from "@/ui/common/utils/maxDecimals";
@@ -30,17 +29,6 @@ export function StakingFeesSection() {
   );
 
   const { calculateFeeAmount } = useStakingService();
-  const { stakingInfo } = useStakingState();
-
-  useEffect(() => {
-    if (stakingInfo?.defaultFeeRate !== undefined) {
-      setValue("feeRate", stakingInfo.defaultFeeRate.toString(), {
-        shouldValidate: true,
-        shouldDirty: true,
-        shouldTouch: true,
-      });
-    }
-  }, [stakingInfo?.defaultFeeRate, setValue]);
 
   useEffect(() => {
     let cancelled = false;
@@ -162,6 +150,7 @@ export function StakingFeesSection() {
         open={feeModalVisible}
         onClose={() => setFeeModalVisible(false)}
         onSubmit={handleFeeRateSubmit}
+        currentFeeRate={Number(feeRate || 0)}
       />
     </>
   );

services/simple-staking/src/ui/common/components/Staking/FeeModal/index.tsx

@@ -22,9 +22,15 @@ interface FeeModalProps {
   open?: boolean;
   onSubmit?: (value: number) => void;
   onClose?: () => void;
+  currentFeeRate?: number;
 }
 
-export function FeeModal({ open, onSubmit, onClose }: FeeModalProps) {
+export function FeeModal({
+  open,
+  onSubmit,
+  onClose,
+  currentFeeRate,
+}: FeeModalProps) {
   const [selectedValue, setSelectedValue] = useState("");
   const [customFee, setCustomFee] = useState("");
   const customFeeRef = useRef<HTMLInputElement>(null);
@@ -45,6 +51,32 @@ export function FeeModal({ open, onSubmit, onClose }: FeeModalProps) {
     }
   }, [selectedValue]);
 
+  // Initialize selection based on current fee rate when opening
+  useEffect(() => {
+    if (!open || isLoading) return;
+
+    const fee = Number(currentFeeRate);
+    if (!fee || !Number.isFinite(fee)) {
+      setSelectedValue("");
+      setCustomFee("");
+      return;
+    }
+
+    if (fee === fastestFee) {
+      setSelectedValue("fast");
+      setCustomFee("");
+    } else if (fee === mediumFee) {
+      setSelectedValue("medium");
+      setCustomFee("");
+    } else if (fee === lowestFee) {
+      setSelectedValue("slow");
+      setCustomFee("");
+    } else {
+      setSelectedValue("custom");
+      setCustomFee(fee.toString());
+    }
+  }, [open, isLoading, currentFeeRate, fastestFee, mediumFee, lowestFee]);
+
   const feeOptions = [
     {
       label: (

services/simple-staking/src/ui/common/components/Tabs/Tabs.tsx

@@ -13,6 +13,7 @@ interface TabsProps {
   activeTab?: string;
   onTabChange?: (tabId: string) => void;
   className?: string;
+  keepMounted?: boolean;
 }
 
 export const Tabs = ({
@@ -21,6 +22,7 @@ export const Tabs = ({
   activeTab: controlledActiveTab,
   onTabChange,
   className,
+  keepMounted,
 }: TabsProps) => {
   const [internalActiveTab, setInternalActiveTab] = useState(
     defaultActiveTab || items[0]?.id || "",
@@ -69,14 +71,30 @@ export const Tabs = ({
         ))}
       </div>
 
-      <div
-        className="mt-6 min-h-[450px]"
-        role="tabpanel"
-        id={`panel-${activeTab}`}
-        aria-labelledby={`tab-${activeTab}`}
-      >
-        {activeContent}
-      </div>
+      {keepMounted ? (
+        <div className="mt-6 min-h-[450px]">
+          {items.map((item) => (
+            <div
+              key={item.id}
+              role="tabpanel"
+              id={`panel-${item.id}`}
+              aria-labelledby={`tab-${item.id}`}
+              className={twMerge(activeTab === item.id ? "" : "hidden")}
+            >
+              {item.content}
+            </div>
+          ))}
+        </div>
+      ) : (
+        <div
+          className="mt-6 min-h-[450px]"
+          role="tabpanel"
+          id={`panel-${activeTab}`}
+          aria-labelledby={`tab-${activeTab}`}
+        >
+          {activeContent}
+        </div>
+      )}
     </div>
   );
 };

services/simple-staking/src/ui/common/page.tsx

@@ -68,6 +68,7 @@ const Home = () => {
         defaultActiveTab="stake"
         activeTab={activeTab}
         onTabChange={setActiveTab}
+        keepMounted
       />
     </Container>
   );

services/simple-staking/src/ui/common/state/FormPersistenceState.tsx

@@ -0,0 +1,43 @@
+import { useMemo, useState, type PropsWithChildren } from "react";
+
+import { type MultistakingFormFields } from "@/ui/common/state/MultistakingState";
+import { createStateUtils } from "@/ui/common/utils/createStateUtils";
+import { type StakingFormFields } from "@/ui/baby/widgets/StakingForm";
+
+interface FormPersistenceContext {
+  btcStakeDraft?: Partial<MultistakingFormFields> | undefined;
+  babyStakeDraft?: Partial<StakingFormFields> | undefined;
+  setBtcStakeDraft: (draft?: Partial<MultistakingFormFields>) => void;
+  setBabyStakeDraft: (draft?: Partial<StakingFormFields>) => void;
+}
+
+const { StateProvider, useState: useFormPersistenceState } =
+  createStateUtils<FormPersistenceContext>({
+    btcStakeDraft: undefined,
+    babyStakeDraft: undefined,
+    setBtcStakeDraft: () => {},
+    setBabyStakeDraft: () => {},
+  });
+
+export function FormPersistenceState({ children }: PropsWithChildren) {
+  const [btcStakeDraft, setBtcStakeDraft] = useState<
+    Partial<MultistakingFormFields> | undefined
+  >(undefined);
+  const [babyStakeDraft, setBabyStakeDraft] = useState<
+    Partial<StakingFormFields> | undefined
+  >(undefined);
+
+  const context = useMemo(
+    () => ({
+      btcStakeDraft,
+      babyStakeDraft,
+      setBtcStakeDraft,
+      setBabyStakeDraft,
+    }),
+    [btcStakeDraft, babyStakeDraft, setBtcStakeDraft, setBabyStakeDraft],
+  );
+
+  return <StateProvider value={context}>{children}</StateProvider>;
+}
+
+export { useFormPersistenceState };

services/simple-staking/src/ui/common/state/index.tsx

@@ -18,12 +18,14 @@ import { BalanceState } from "./BalanceState";
 import { DelegationState } from "./DelegationState";
 import { DelegationV2State } from "./DelegationV2State";
 import { FinalityProviderState } from "./FinalityProviderState";
+import { FormPersistenceState } from "./FormPersistenceState";
 import { RewardsState } from "./RewardState";
 import { StakingExpansionState } from "./StakingExpansionState";
 import { StakingState } from "./StakingState";
 
 // The order of the states is important for the state provider
 const STATE_LIST = [
+  FormPersistenceState,
   DelegationState,
   DelegationV2State,
   FinalityProviderState,