From 3e19ed04cf03743067832acdabd89bd1b1d422ea Mon Sep 17 00:00:00 2001
From: antoncoding <antonassocareer@gmail.com>
Date: Sun, 30 Nov 2025 11:27:01 +0800
Subject: [PATCH] fix: statle state on previews

---
 src/components/SupplyModalContent.tsx | 28 ++++++++++++-------
 src/components/SupplyModalV2.tsx      | 39 ++++++++++++++++++++-------
 src/hooks/useSupplyMarket.ts          | 11 ++++++--
 3 files changed, 57 insertions(+), 21 deletions(-)

diff --git a/src/components/SupplyModalContent.tsx b/src/components/SupplyModalContent.tsx
index f2d1a7be..94d812a7 100644
--- a/src/components/SupplyModalContent.tsx
+++ b/src/components/SupplyModalContent.tsx
@@ -1,5 +1,6 @@
 import React, { useCallback } from 'react';
 import { Switch } from '@heroui/react';
+import { ReloadIcon } from '@radix-ui/react-icons';
 import { useAccount } from 'wagmi';
 import Input from '@/components/Input/Input';
 import AccountConnect from '@/components/layout/header/AccountConnect';
@@ -53,10 +54,9 @@ export function SupplyModalContent({
     supplyPending,
     approveAndSupply,
     signAndSupply,
+    refetch: refetchBalance
   } = useSupplyMarket(market, onSuccess);
 
-  console.log('tokenBalance', tokenBalance)
-
   // Handle supply amount change
   const handleSupplyAmountChange = useCallback(
     (amount: bigint) => {
@@ -111,13 +111,23 @@ export function SupplyModalContent({
                 <div>
                   <div className="flex items-center justify-between">
                     <span className="opacity-80">Supply amount</span>
-                    <p className="font-inter text-xs opacity-50">
-                      Balance:{' '}
-                      {useEth
-                        ? formatBalance(ethBalance ?? BigInt(0), 18)
-                        : formatBalance(tokenBalance ?? BigInt(0), market.loanAsset.decimals)}{' '}
-                      {useEth ? getNativeTokenSymbol(market.morphoBlue.chain.id) : market.loanAsset.symbol}
-                    </p>
+                    <div className="flex items-center gap-1.5">
+                      <p className="font-inter text-xs opacity-50">
+                        Balance:{' '}
+                        {useEth
+                          ? formatBalance(ethBalance ?? BigInt(0), 18)
+                          : formatBalance(tokenBalance ?? BigInt(0), market.loanAsset.decimals)}{' '}
+                        {useEth ? getNativeTokenSymbol(market.morphoBlue.chain.id) : market.loanAsset.symbol}
+                      </p>
+                      <button
+                        type="button"
+                        onClick={() => void refetchBalance()}
+                        className="opacity-50 transition hover:opacity-100"
+                        aria-label="Refetch balance"
+                      >
+                        <ReloadIcon className="h-3 w-3" />
+                      </button>
+                    </div>
                   </div>
 
                   <div className="mt-2 flex items-start justify-between">
diff --git a/src/components/SupplyModalV2.tsx b/src/components/SupplyModalV2.tsx
index e12e4117..8dab762e 100644
--- a/src/components/SupplyModalV2.tsx
+++ b/src/components/SupplyModalV2.tsx
@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { useEffect, useMemo, useState } from 'react';
 import { LuArrowRightLeft } from "react-icons/lu";
 import { Modal, ModalBody, ModalHeader } from '@/components/common/Modal';
 import { useFreshMarketsState } from '@/hooks/useFreshMarketsState';
@@ -28,12 +28,31 @@ export function SupplyModalV2({
   const [supplyPreviewAmount, setSupplyPreviewAmount] = useState<bigint | undefined>();
   const [withdrawPreviewAmount, setWithdrawPreviewAmount] = useState<bigint | undefined>();
 
+  // Reset preview amounts and mode when modal is first opened to prevent stale state
+  useEffect(() => {
+    setMode(defaultMode);
+    setSupplyPreviewAmount(undefined);
+    setWithdrawPreviewAmount(undefined);
+  }, [defaultMode]);
+
   // Fetch fresh market state from RPC to avoid stale liquidity/supply data
   const { markets: freshMarkets } = useFreshMarketsState([market]);
   const activeMarket = freshMarkets?.[0] ?? market;
 
   const hasPosition = position && BigInt(position.state.supplyAssets) > 0n;
 
+  // Calculate supply delta for preview based on current mode and amounts
+  // Only use positive values to prevent incorrect APY direction
+  const supplyDelta = useMemo(() => {
+    if (mode === 'supply') {
+      // Supply mode: positive delta if amount is valid
+      return supplyPreviewAmount && supplyPreviewAmount > 0n ? supplyPreviewAmount : undefined;
+    } else {
+      // Withdraw mode: negative delta (withdrawal) if amount is valid
+      return withdrawPreviewAmount && withdrawPreviewAmount > 0n ? -withdrawPreviewAmount : undefined;
+    }
+  }, [mode, supplyPreviewAmount, withdrawPreviewAmount]);
+
   return (
     <Modal
       isOpen
@@ -61,8 +80,14 @@ export function SupplyModalV2({
           hasPosition ? (
             <button
               type="button"
-              onClick={() => setMode(mode === 'supply' ? 'withdraw' : 'supply')}
-              className="flex items-center gap-1 text-sm font-medium text-primary transition hover:text-white"
+              onClick={() => {
+                const newMode = mode === 'supply' ? 'withdraw' : 'supply';
+                setMode(newMode);
+                // Reset preview amounts when switching modes to prevent stale state
+                setSupplyPreviewAmount(undefined);
+                setWithdrawPreviewAmount(undefined);
+              }}
+              className="flex items-center gap-1 text-sm font-medium text-primary transition hover:text-secondary"
             >
               <LuArrowRightLeft className="h-3 w-3 rotate-90" />
               {mode === 'supply' ? 'Withdraw' : 'Supply'}
@@ -77,13 +102,7 @@ export function SupplyModalV2({
           defaultCollapsed
           mode="supply"
           showRewards
-          supplyDelta={
-            mode === 'supply'
-              ? supplyPreviewAmount
-              : withdrawPreviewAmount
-              ? -withdrawPreviewAmount
-              : undefined
-          }
+          supplyDelta={supplyDelta}
         />
 
         {mode === 'supply' ? (
diff --git a/src/hooks/useSupplyMarket.ts b/src/hooks/useSupplyMarket.ts
index e56a2e56..aafe86fe 100644
--- a/src/hooks/useSupplyMarket.ts
+++ b/src/hooks/useSupplyMarket.ts
@@ -40,6 +40,7 @@ export type UseSupplyMarketReturn = {
   // Actions
   approveAndSupply: () => Promise<void>;
   signAndSupply: () => Promise<void>;
+  refetch: () => void;
 };
 
 export function useSupplyMarket(market: Market, onSuccess?: () => void): UseSupplyMarketReturn {
@@ -56,14 +57,14 @@ export function useSupplyMarket(market: Market, onSuccess?: () => void): UseSupp
   const toast = useStyledToast();
 
   // Get token balance
-  const { data: tokenBalance } = useBalance({
+  const { data: tokenBalance, refetch: refetchToken } = useBalance({
     token: market.loanAsset.address as `0x${string}`,
     address: account,
     chainId: market.morphoBlue.chain.id,
   });
 
   // Get ETH balance
-  const { data: ethBalance } = useBalance({
+  const { data: ethBalance, refetch: refetchETH } = useBalance({
     address: account,
     chainId: market.morphoBlue.chain.id,
   });
@@ -92,6 +93,11 @@ export function useSupplyMarket(market: Market, onSuccess?: () => void): UseSupp
     tokenSymbol: market.loanAsset.symbol,
   });
 
+  const refetch = useCallback(() => {
+    void refetchToken()
+    void refetchETH()
+  }, [refetchETH, refetchToken])
+
   // Transaction handler
   const { isConfirming: supplyPending, sendTransactionAsync } = useTransactionWithToast({
     toastId: 'supply',
@@ -348,6 +354,7 @@ export function useSupplyMarket(market: Market, onSuccess?: () => void): UseSupp
     // Balance data
     tokenBalance: tokenBalance?.value,
     ethBalance: ethBalance?.value,
+    refetch: refetch,
 
     // Transaction state
     isApproved,