refactor: save entire data structure on submission

src/components/tx-flow/common/TxLayout.tsx

@@ -1,5 +1,5 @@
 import type { ReactNode } from 'react'
-import { Box, Container, Grid, Paper, Typography } from '@mui/material'
+import { Box, Container, Grid, Paper, Typography, Button } from '@mui/material'
 import type { TransactionSummary } from '@safe-global/safe-gateway-typescript-sdk'
 import { ProgressBar } from '@/components/common/ProgressBar'
 import SafeTxProvider from '../SafeTxProvider'
@@ -11,9 +11,10 @@ type TxLayoutProps = {
   children: ReactNode
   step?: number
   txSummary?: TransactionSummary
+  onBack?: () => void
 }
 
-const TxLayout = ({ title, children, step = 0, txSummary }: TxLayoutProps) => {
+const TxLayout = ({ title, children, step = 0, txSummary, onBack }: TxLayoutProps) => {
   const steps = Array.isArray(children) ? children : [children]
   const progress = Math.round(((step + 1) / steps.length) * 100)
 
@@ -29,17 +30,20 @@ const TxLayout = ({ title, children, step = 0, txSummary }: TxLayoutProps) => {
 
           <Grid item container xs={12} gap={3}>
             <Grid item xs={7} component={Paper}>
-              {<ProgressBar value={progress} />}
+              <ProgressBar value={progress} />
 
               <Box display="flex" justifyContent="flex-end" py={2} px={3}>
                 <TxNonce />
               </Box>
 
-              {steps.map((children, index) => (
-                <div key={index} style={{ display: index === step ? '' : 'none' }}>
-                  {children}
-                </div>
-              ))}
+              <div>
+                {steps[step]}
+                {onBack && step > 0 && (
+                  <Button variant="contained" onClick={onBack}>
+                    Back
+                  </Button>
+                )}
+              </div>
             </Grid>
 
             <Grid item xs={4}>

src/components/tx-flow/flows/NftTransfer/ReviewNftBatch.tsx

@@ -12,7 +12,6 @@ import { SafeTxContext } from '../../SafeTxProvider'
 type ReviewNftBatchProps = {
   params: SubmittedNftTransferParams
   onSubmit: () => void
-  onBack: () => void
   txNonce?: number
 }
 

src/components/tx-flow/flows/NftTransfer/index.tsx

@@ -26,13 +26,13 @@ const NftTransferFlow = ({ txNonce, ...params }: NftTransferFlowProps) => {
   ])
 
   const steps = [
-    <SendNftBatch key={0} params={data[0]} onSubmit={(formData) => nextStep<1>(formData)} />,
+    <SendNftBatch key={0} params={data[0]} onSubmit={(formData) => nextStep([formData, formData])} />,
 
-    data[1] && <ReviewNftBatch key={1} params={data[1]} txNonce={txNonce} onSubmit={() => null} onBack={prevStep} />,
+    data[1] && <ReviewNftBatch key={1} params={data[1]} txNonce={txNonce} onSubmit={() => null} />,
   ]
 
   return (
-    <TxLayout title="Send tokens" step={step}>
+    <TxLayout title="Send tokens" step={step} onBack={prevStep}>
       {steps}
     </TxLayout>
   )

src/components/tx-flow/flows/TokenTransfer/CreateTokenTransfer.tsx

@@ -33,13 +33,8 @@ import SpendingLimitRow from '@/components/tx/SpendingLimitRow'
 import NumberField from '@/components/common/NumberField'
 import InputValueHelper from '@/components/common/InputValueHelper'
 import { validateDecimalLength, validateLimitedAmount } from '@/utils/validation'
-import {
-  AutocompleteItem,
-  SendAssetsField,
-  type SendAssetsFormData,
-  SendTxType,
-} from '@/components/tx/modals/TokenTransferModal/SendAssetsForm'
-import type { TokenTransferParams } from '.'
+import { AutocompleteItem } from '@/components/tx/modals/TokenTransferModal/SendAssetsForm'
+import { type TokenTransferParams, TokenTransferFields, TokenTransferType } from '.'
 
 const CreateTokenTransfer = ({
   params,
@@ -60,15 +55,13 @@ const CreateTokenTransfer = ({
   const spendingLimits = useAppSelector(selectSpendingLimits)
   const wallet = useWallet()
 
-  const formMethods = useForm<SendAssetsFormData>({
+  const formMethods = useForm<TokenTransferParams>({
     defaultValues: {
-      [SendAssetsField.recipient]: params.recipient || '',
-      [SendAssetsField.tokenAddress]: params.tokenAddress || '',
-      [SendAssetsField.amount]: params.amount || '',
-      [SendAssetsField.type]: disableSpendingLimit
-        ? SendTxType.multiSig
+      ...params,
+      [TokenTransferFields.type]: disableSpendingLimit
+        ? TokenTransferType.multiSig
         : isOnlySpendingLimitBeneficiary
-        ? SendTxType.spendingLimit
+        ? TokenTransferType.spendingLimit
         : params.type,
     },
     mode: 'onChange',
@@ -85,17 +78,17 @@ const CreateTokenTransfer = ({
     control,
   } = formMethods
 
-  const recipient = watch(SendAssetsField.recipient)
+  const recipient = watch(TokenTransferFields.recipient)
 
   // Selected token
-  const tokenAddress = watch(SendAssetsField.tokenAddress)
+  const tokenAddress = watch(TokenTransferFields.tokenAddress)
   const selectedToken = tokenAddress
     ? balances.items.find((item) => item.tokenInfo.address === tokenAddress)
     : undefined
 
-  const type = watch(SendAssetsField.type)
+  const type = watch(TokenTransferFields.type)
   const spendingLimit = useSpendingLimit(selectedToken?.tokenInfo)
-  const isSpendingLimitType = type === SendTxType.spendingLimit
+  const isSpendingLimitType = type === TokenTransferType.spendingLimit
   const spendingLimitAmount = spendingLimit ? BigNumber.from(spendingLimit.amount).sub(spendingLimit.spent) : undefined
   const totalAmount = BigNumber.from(selectedToken?.balance || 0)
   const maxAmount = isSpendingLimitType
@@ -122,7 +115,7 @@ const CreateTokenTransfer = ({
         ? spendingLimitAmount.toString()
         : selectedToken.balance
 
-    setValue(SendAssetsField.amount, safeFormatUnits(amount, selectedToken.tokenInfo.decimals), {
+    setValue(TokenTransferFields.amount, safeFormatUnits(amount, selectedToken.tokenInfo.decimals), {
       shouldValidate: true,
     })
   }, [isSpendingLimitType, selectedToken, setValue, spendingLimitAmount])
@@ -138,16 +131,16 @@ const CreateTokenTransfer = ({
 
           <FormControl fullWidth sx={{ mb: 2, mt: 1 }}>
             {addressBook[recipient] ? (
-              <Box onClick={() => setValue(SendAssetsField.recipient, '')}>
+              <Box onClick={() => setValue(TokenTransferFields.recipient, '')}>
                 <SendToBlock address={recipient} />
               </Box>
             ) : (
-              <AddressBookInput name={SendAssetsField.recipient} label="Recipient" />
+              <AddressBookInput name={TokenTransferFields.recipient} label="Recipient" />
             )}
           </FormControl>
 
           <Controller
-            name={SendAssetsField.tokenAddress}
+            name={TokenTransferFields.tokenAddress}
             control={control}
             rules={{ required: true }}
             render={({ fieldState, field }) => (
@@ -162,7 +155,7 @@ const CreateTokenTransfer = ({
                   {...field}
                   onChange={(e) => {
                     field.onChange(e)
-                    resetField(SendAssetsField.amount)
+                    resetField(TokenTransferFields.amount)
                   }}
                 >
                   {balancesItems.map((item) => (
@@ -203,10 +196,10 @@ const CreateTokenTransfer = ({
               }}
               // @see https://github.com/react-hook-form/react-hook-form/issues/220
               InputLabelProps={{
-                shrink: !!watch(SendAssetsField.amount),
+                shrink: !!watch(TokenTransferFields.amount),
               }}
               required
-              {...register(SendAssetsField.amount, {
+              {...register(TokenTransferFields.amount, {
                 required: true,
                 validate: (val) => {
                   const decimals = selectedToken?.tokenInfo.decimals

src/components/tx-flow/flows/TokenTransfer/ReviewTokenTransfer.tsx

@@ -17,7 +17,6 @@ const ReviewTokenTransfer = ({
 }: {
   params: TokenTransferParams
   onSubmit: () => void
-  onBack: () => void
   txNonce?: number
 }) => {
   const { setSafeTx, setSafeTxError, setNonce } = useContext(SafeTxContext)

src/components/tx-flow/flows/TokenTransfer/index.tsx

@@ -1,35 +1,61 @@
+import merge from 'lodash/merge'
 import TxLayout from '@/components/tx-flow/common/TxLayout'
 import useTxStepper from '../../useTxStepper'
 import CreateTokenTransfer from './CreateTokenTransfer'
 import ReviewTokenTransfer from './ReviewTokenTransfer'
+import { useMemo } from 'react'
 
-export enum SendTxType {
+export enum TokenTransferType {
   multiSig = 'multiSig',
   spendingLimit = 'spendingLimit',
 }
 
+export enum TokenTransferFields {
+  recipient = 'recipient',
+  tokenAddress = 'tokenAddress',
+  amount = 'amount',
+  type = 'type',
+}
+
 export type TokenTransferParams = {
-  recipient?: string
-  tokenAddress?: string
-  amount?: string
-  type?: SendTxType
+  [TokenTransferFields.recipient]: string
+  [TokenTransferFields.tokenAddress]: string
+  [TokenTransferFields.amount]: string
+  [TokenTransferFields.type]: TokenTransferType
 }
 
-type TokenTransferFlowProps = TokenTransferParams & {
+type TokenTransferFlowProps = Partial<TokenTransferParams> & {
   txNonce?: number
 }
 
+const defaultData: TokenTransferParams = {
+  recipient: '',
+  tokenAddress: '',
+  amount: '',
+  type: TokenTransferType.multiSig,
+}
+
 const TokenTransferFlow = ({ txNonce, ...params }: TokenTransferFlowProps) => {
-  const { data, step, nextStep, prevStep } = useTxStepper<[TokenTransferParams, TokenTransferParams]>([params, params])
+  const initialData = useMemo(() => merge({}, defaultData, params), [params])
+
+  const { data, step, nextStep, prevStep } = useTxStepper<[TokenTransferParams, TokenTransferParams | undefined]>([
+    initialData,
+    undefined,
+  ])
 
   const steps = [
-    <CreateTokenTransfer key={0} params={data[0]} txNonce={txNonce} onSubmit={(formData) => nextStep<1>(formData)} />,
+    <CreateTokenTransfer
+      key={0}
+      params={data[0]}
+      txNonce={txNonce}
+      onSubmit={(formData) => nextStep([formData, formData])}
+    />,
 
-    <ReviewTokenTransfer key={1} txNonce={txNonce} params={data[1]} onSubmit={() => null} onBack={prevStep} />,
+    data[1] && <ReviewTokenTransfer key={1} params={data[1]} txNonce={txNonce} onSubmit={() => null} />,
   ]
 
   return (
-    <TxLayout title="Send tokens" step={step}>
+    <TxLayout title="Send tokens" step={step} onBack={prevStep}>
       {steps}
     </TxLayout>
   )

src/components/tx-flow/useTxStepper.tsx

@@ -1,15 +1,12 @@
 import { useCallback, useState } from 'react'
 
-const useTxStepper = <T extends Array<any>>(initialData: T) => {
+const useTxStepper = <T extends Array<unknown>>(initialData: T) => {
   const [step, setStep] = useState(0)
   const [data, setData] = useState<T>(initialData)
 
-  const nextStep = useCallback(<S extends keyof T>(stepData: T[S]) => {
-    setStep((prevStep) => {
-      const newStep = prevStep + 1
-      setData((prevData) => ({ ...prevData, [newStep]: stepData }))
-      return newStep
-    })
+  const nextStep = useCallback((entireData: T) => {
+    setData(entireData)
+    setStep((prevStep) => prevStep + 1)
   }, [])
 
   const prevStep = useCallback(() => {

src/components/tx/SignOrExecuteForm/SignForm.tsx

@@ -1,4 +1,4 @@
-import { type ReactElement, type SyntheticEvent, useState } from 'react'
+import { type ReactElement, type SyntheticEvent, useContext, useState } from 'react'
 import { Button, Typography } from '@mui/material'
 
 import ErrorMessage from '@/components/tx/ErrorMessage'
@@ -8,6 +8,7 @@ import CheckWallet from '@/components/common/CheckWallet'
 import { useTxActions } from './hooks'
 import type { SignOrExecuteProps } from '.'
 import type { SafeTransaction } from '@safe-global/safe-core-sdk-types'
+import { TxModalContext } from '@/components/tx-flow'
 
 const SignForm = ({
   safeTx,
@@ -29,6 +30,7 @@ const SignForm = ({
   // Hooks
   const isOwner = useIsSafeOwner()
   const { signTx } = useTxActions()
+  const { setTxFlow } = useContext(TxModalContext)
 
   // Check that the transaction is executable
   const isCreation = !txId
@@ -41,6 +43,7 @@ const SignForm = ({
 
     try {
       await signTx(safeTx, txId, origin)
+      setTxFlow(undefined)
     } catch (err) {
       logError(Errors._804, (err as Error).message)
       setIsSubmittable(true)