fix: state refresh after action bug

mrgnlabs · Dec 17, 2024 · 6a011cf · 6a011cf
1 parent 81b5165
commit 6a011cf
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 43 deletions.
diff --git a/apps/marginfi-v2-trading/src/components/common/trade-box-v2/hooks/use-action-amounts.ts b/apps/marginfi-v2-trading/src/components/common/trade-box-v2/hooks/use-action-amounts.ts
@@ -1,44 +1,49 @@
 import React from "react";
 
-import { ActionType, ExtendedBankInfo } from "@mrgnlabs/marginfi-v2-ui-state";
 import { WSOL_MINT } from "@mrgnlabs/mrgn-common";
 
 import { useAmountDebounce } from "~/hooks/useAmountDebounce";
-import { ArenaBank, ArenaPoolV2Extended } from "~/types/trade-store.types";
+import { ArenaPoolV2Extended } from "~/types/trade-store.types";
+import { PublicKey } from "@solana/web3.js";
 
 export function useActionAmounts({
   amountRaw,
   activePool,
-  collateralBank,
+  selectedBankPk,
   nativeSolBalance,
 }: {
   amountRaw: string;
   activePool: ArenaPoolV2Extended | null;
-  collateralBank: ArenaBank | null;
+  selectedBankPk: PublicKey | null;
   nativeSolBalance: number;
 }) {
   const amount = React.useMemo(() => {
     const strippedAmount = amountRaw.replace(/,/g, "");
     return isNaN(Number.parseFloat(strippedAmount)) ? 0 : Number.parseFloat(strippedAmount);
   }, [amountRaw]);
 
+  const bank = React.useMemo(() => {
+    if (!selectedBankPk) return null;
+    return [activePool?.tokenBank, activePool?.quoteBank].find((bank) => bank?.address.equals(selectedBankPk));
+  }, [selectedBankPk, activePool]);
+
   const debouncedAmount = useAmountDebounce<number | null>(amount, 500);
 
   const walletAmount = React.useMemo(
     () =>
-      collateralBank?.info.state.mint?.equals && collateralBank?.info.state.mint?.equals(WSOL_MINT)
-        ? collateralBank?.userInfo.tokenAccount.balance + nativeSolBalance
-        : collateralBank?.userInfo.tokenAccount.balance,
-    [nativeSolBalance, collateralBank]
+      bank?.info.state.mint?.equals && bank?.info.state.mint?.equals(WSOL_MINT)
+        ? bank?.userInfo.tokenAccount.balance + nativeSolBalance
+        : bank?.userInfo.tokenAccount.balance,
+    [nativeSolBalance, bank]
   );
 
   const maxAmount = React.useMemo(() => {
-    if (!collateralBank) {
+    if (!bank) {
       return 0;
     }
 
-    return collateralBank.userInfo.maxDeposit;
-  }, [collateralBank]);
+    return bank.userInfo.maxDeposit;
+  }, [bank]);
 
   return {
     amount,

diff --git a/apps/marginfi-v2-trading/src/components/common/trade-box-v2/store/trade-box-store.tsx b/apps/marginfi-v2-trading/src/components/common/trade-box-v2/store/trade-box-store.tsx
@@ -7,6 +7,7 @@ import { SimulationResult } from "@mrgnlabs/marginfi-client-v2";
 import { ArenaBank } from "~/types/trade-store.types";
 
 import { TradeSide } from "..";
+import { PublicKey } from "@solana/web3.js";
 
 interface TradeBoxState {
   // State
@@ -15,8 +16,8 @@ interface TradeBoxState {
   leverage: number;
   maxLeverage: number;
 
-  selectedBank: ArenaBank | null;
-  selectedSecondaryBank: ArenaBank | null;
+  selectedBankPk: PublicKey | null;
+  selectedSecondaryBankPk: PublicKey | null;
 
   simulationResult: SimulationResult | null;
   actionTxns: TradeActionTxns;
@@ -32,8 +33,8 @@ interface TradeBoxState {
   setSimulationResult: (result: SimulationResult | null) => void;
   setActionTxns: (actionTxns: TradeActionTxns) => void;
   setErrorMessage: (errorMessage: ActionMessageType | null) => void;
-  setSelectedBank: (bank: ArenaBank | null) => void;
-  setSelectedSecondaryBank: (bank: ArenaBank | null) => void;
+  setSelectedBankPk: (bankAddress: PublicKey | null) => void;
+  setSelectedSecondaryBankPk: (bankAddress: PublicKey | null) => void;
   setMaxLeverage: (maxLeverage: number) => void;
 }
 
@@ -43,8 +44,8 @@ const initialState = {
   leverage: 0,
   simulationResult: null,
   errorMessage: null,
-  selectedBank: null,
-  selectedSecondaryBank: null,
+  selectedBankPk: null,
+  selectedSecondaryBankPk: null,
   maxLeverage: 0,
 
   actionTxns: {
@@ -122,13 +123,13 @@ const stateCreator: StateCreator<TradeBoxState, [], []> = (set, get) => ({
     set({ errorMessage: errorMessage });
   },
 
-  setSelectedBank(tokenBank) {
-    const selectedBank = get().selectedBank;
-    const hasBankChanged = !tokenBank || !selectedBank || !tokenBank.address.equals(selectedBank.address);
+  setSelectedBankPk(bankAddress) {
+    const selectedBankPk = get().selectedBankPk;
+    const hasBankChanged = !bankAddress || !selectedBankPk || !bankAddress.equals(selectedBankPk);
 
     if (hasBankChanged) {
       set({
-        selectedBank: tokenBank,
+        selectedBankPk: bankAddress,
         amountRaw: initialState.amountRaw,
         leverage: initialState.leverage,
         actionTxns: initialState.actionTxns,
@@ -137,21 +138,20 @@ const stateCreator: StateCreator<TradeBoxState, [], []> = (set, get) => ({
     }
   },
 
-  setSelectedSecondaryBank(secondaryBank) {
-    const selectedSecondaryBank = get().selectedSecondaryBank;
-    const hasBankChanged =
-      !secondaryBank || !selectedSecondaryBank || !secondaryBank.address.equals(selectedSecondaryBank.address);
+  setSelectedSecondaryBankPk(bankAddress) {
+    const selectedSecondaryBankPk = get().selectedSecondaryBankPk;
+    const hasBankChanged = !bankAddress || !selectedSecondaryBankPk || !bankAddress.equals(selectedSecondaryBankPk);
 
     if (hasBankChanged) {
       set({
-        selectedSecondaryBank: secondaryBank,
+        selectedSecondaryBankPk: bankAddress,
         amountRaw: initialState.amountRaw,
         leverage: initialState.leverage,
         actionTxns: initialState.actionTxns,
         errorMessage: null,
       });
     } else {
-      set({ selectedSecondaryBank: secondaryBank });
+      set({ selectedSecondaryBankPk: bankAddress });
     }
   },
 

diff --git a/apps/marginfi-v2-trading/src/components/common/trade-box-v2/trade-box.tsx b/apps/marginfi-v2-trading/src/components/common/trade-box-v2/trade-box.tsx
@@ -58,8 +58,8 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
     simulationResult,
     actionTxns,
     errorMessage,
-    selectedBank,
-    selectedSecondaryBank,
+    selectedBankPk,
+    selectedSecondaryBankPk,
     maxLeverage,
     refreshState,
     setAmountRaw,
@@ -68,8 +68,8 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
     setSimulationResult,
     setActionTxns,
     setErrorMessage,
-    setSelectedBank,
-    setSelectedSecondaryBank,
+    setSelectedBankPk,
+    setSelectedSecondaryBankPk,
     setMaxLeverage,
   ] = useTradeBoxStore((state) => [
     state.amountRaw,
@@ -78,8 +78,8 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
     state.simulationResult,
     state.actionTxns,
     state.errorMessage,
-    state.selectedBank,
-    state.selectedSecondaryBank,
+    state.selectedBankPk,
+    state.selectedSecondaryBankPk,
     state.maxLeverage,
     state.refreshState,
     state.setAmountRaw,
@@ -88,8 +88,8 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
     state.setSimulationResult,
     state.setActionTxns,
     state.setErrorMessage,
-    state.setSelectedBank,
-    state.setSelectedSecondaryBank,
+    state.setSelectedBankPk,
+    state.setSelectedSecondaryBankPk,
     state.setMaxLeverage,
   ]);
   const [
@@ -129,10 +129,26 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
   const { amount, debouncedAmount, maxAmount } = useActionAmounts({
     amountRaw,
     activePool: activePoolExtended,
-    collateralBank: selectedBank,
+    selectedBankPk,
     nativeSolBalance,
   });
   const debouncedLeverage = useAmountDebounce<number>(leverage, 500);
+  const selectedBank = React.useMemo(() => {
+    if (!selectedBankPk) return null;
+    return (
+      [activePoolExtended.tokenBank, activePoolExtended.quoteBank].find((bank) =>
+        bank?.address.equals(selectedBankPk)
+      ) ?? null
+    );
+  }, [selectedBankPk, activePoolExtended]);
+  const selectedSecondaryBank = React.useMemo(() => {
+    if (!selectedSecondaryBankPk) return null;
+    return (
+      [activePoolExtended.tokenBank, activePoolExtended.quoteBank].find((bank) =>
+        bank?.address.equals(selectedSecondaryBankPk)
+      ) ?? null
+    );
+  }, [selectedSecondaryBankPk, activePoolExtended]);
 
   // States
   const [dynamicActionMessages, setDynamicActionMessages] = React.useState<ActionMessageType[]>([]);
@@ -190,14 +206,14 @@ export const TradeBoxV2 = ({ activePool, side = "long" }: TradeBoxV2Props) => {
   React.useEffect(() => {
     if (activePoolExtended) {
       if (tradeState === "short") {
-        setSelectedBank(activePoolExtended.quoteBank);
-        setSelectedSecondaryBank(activePoolExtended.tokenBank);
+        setSelectedBankPk(activePoolExtended.quoteBank.address);
+        setSelectedSecondaryBankPk(activePoolExtended.tokenBank.address);
       } else {
-        setSelectedBank(activePoolExtended.tokenBank);
-        setSelectedSecondaryBank(activePoolExtended.quoteBank);
+        setSelectedBankPk(activePoolExtended.tokenBank.address);
+        setSelectedSecondaryBankPk(activePoolExtended.quoteBank.address);
       }
     }
-  }, [activePoolExtended, setSelectedBank, setSelectedSecondaryBank, tradeState]);
+  }, [activePoolExtended, setSelectedBankPk, setSelectedSecondaryBankPk, tradeState]);
 
   React.useEffect(() => {
     if (errorMessage && errorMessage.description) {

diff --git a/apps/marginfi-v2-trading/src/store/tradeStoreV2.ts b/apps/marginfi-v2-trading/src/store/tradeStoreV2.ts
@@ -564,8 +564,10 @@ const stateCreator: StateCreator<TradeStoreV2State, [], []> = (set, get) => ({
       tokenAccountMap: storeTokenAccountMap,
     } = get();
 
+    const newStoreBanksByBankPk = { ...storeBanksByBankPk };
+
     extendedBankInfos.map((bank) => {
-      storeBanksByBankPk[bank.address.toBase58()] = bank;
+      newStoreBanksByBankPk[bank.address.toBase58()] = bank;
     });
     if (marginfiAccount) {
       storeMarginfiAccountByGroupPk[args.groupPk.toBase58()] = marginfiAccount;
@@ -577,7 +579,7 @@ const stateCreator: StateCreator<TradeStoreV2State, [], []> = (set, get) => ({
     set({
       nativeSolBalance,
       marginfiAccountByGroupPk: storeMarginfiAccountByGroupPk,
-      banksByBankPk: storeBanksByBankPk,
+      banksByBankPk: newStoreBanksByBankPk,
       tokenAccountMap: storeTokenAccountMap,
     });
   },