From 479fabf474cf6fc9ec683009c13b0842a7fb0582 Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 08:58:02 +0100
Subject: [PATCH 1/6] feat: add updated contracts on testnet

---
 src/abis/umbrellaBatchHelper.ts           |  50 ++--
 src/abis/umbrellaDataAggregationHelper.ts | 268 +++++++++++-----------
 src/constants/markets.ts                  |   6 +-
 3 files changed, 171 insertions(+), 153 deletions(-)

diff --git a/src/abis/umbrellaBatchHelper.ts b/src/abis/umbrellaBatchHelper.ts
index f642aa8..8ab1cd8 100644
--- a/src/abis/umbrellaBatchHelper.ts
+++ b/src/abis/umbrellaBatchHelper.ts
@@ -1,33 +1,26 @@
 export const UMBRELLA_BATCH_HELPER_ABI = [
   {
     inputs: [
-      { internalType: "address", name: "rewardsController", type: "address" },
-      { internalType: "address", name: "owner", type: "address" },
+      { internalType: "address", name: "rewardsController_", type: "address" },
+      { internalType: "address", name: "owner_", type: "address" },
     ],
     stateMutability: "nonpayable",
     type: "constructor",
   },
-  {
-    inputs: [{ internalType: "address", name: "target", type: "address" }],
-    name: "AddressEmptyCode",
-    type: "error",
-  },
+  { inputs: [{ internalType: "address", name: "target", type: "address" }], name: "AddressEmptyCode", type: "error" },
+  { inputs: [], name: "EnforcedPause", type: "error" },
   { inputs: [], name: "EthTransferFailed", type: "error" },
+  { inputs: [], name: "ExpectedPause", type: "error" },
   { inputs: [], name: "FailedCall", type: "error" },
   { inputs: [], name: "InvalidEdgeToken", type: "error" },
   { inputs: [], name: "NotInitializedStake", type: "error" },
   { inputs: [], name: "OnlyRescueGuardian", type: "error" },
-  {
-    inputs: [{ internalType: "address", name: "owner", type: "address" }],
-    name: "OwnableInvalidOwner",
-    type: "error",
-  },
+  { inputs: [{ internalType: "address", name: "owner", type: "address" }], name: "OwnableInvalidOwner", type: "error" },
   {
     inputs: [{ internalType: "address", name: "account", type: "address" }],
     name: "OwnableUnauthorizedAccount",
     type: "error",
   },
-  { inputs: [], name: "PathFromOnlyStakeTokenCouldBeInitialized", type: "error" },
   {
     inputs: [{ internalType: "address", name: "token", type: "address" }],
     name: "SafeERC20FailedOperation",
@@ -37,7 +30,7 @@ export const UMBRELLA_BATCH_HELPER_ABI = [
   { inputs: [], name: "ZeroAmount", type: "error" },
   {
     anonymous: false,
-    inputs: [{ indexed: false, internalType: "address", name: "stakeToken", type: "address" }],
+    inputs: [{ indexed: true, internalType: "address", name: "stakeToken", type: "address" }],
     name: "AssetPathInitialized",
     type: "event",
   },
@@ -71,6 +64,18 @@ export const UMBRELLA_BATCH_HELPER_ABI = [
     name: "OwnershipTransferred",
     type: "event",
   },
+  {
+    anonymous: false,
+    inputs: [{ indexed: false, internalType: "address", name: "account", type: "address" }],
+    name: "Paused",
+    type: "event",
+  },
+  {
+    anonymous: false,
+    inputs: [{ indexed: false, internalType: "address", name: "account", type: "address" }],
+    name: "Unpaused",
+    type: "event",
+  },
   {
     inputs: [],
     name: "REWARDS_CONTROLLER",
@@ -187,6 +192,14 @@ export const UMBRELLA_BATCH_HELPER_ABI = [
     stateMutability: "view",
     type: "function",
   },
+  { inputs: [], name: "pause", outputs: [], stateMutability: "nonpayable", type: "function" },
+  {
+    inputs: [],
+    name: "paused",
+    outputs: [{ internalType: "bool", name: "", type: "bool" }],
+    stateMutability: "view",
+    type: "function",
+  },
   {
     inputs: [
       {
@@ -226,13 +239,7 @@ export const UMBRELLA_BATCH_HELPER_ABI = [
     stateMutability: "nonpayable",
     type: "function",
   },
-  {
-    inputs: [],
-    name: "renounceOwnership",
-    outputs: [],
-    stateMutability: "nonpayable",
-    type: "function",
-  },
+  { inputs: [], name: "renounceOwnership", outputs: [], stateMutability: "nonpayable", type: "function" },
   {
     inputs: [{ internalType: "address", name: "newOwner", type: "address" }],
     name: "transferOwnership",
@@ -240,6 +247,7 @@ export const UMBRELLA_BATCH_HELPER_ABI = [
     stateMutability: "nonpayable",
     type: "function",
   },
+  { inputs: [], name: "unpause", outputs: [], stateMutability: "nonpayable", type: "function" },
   {
     inputs: [],
     name: "whoCanRescue",
diff --git a/src/abis/umbrellaDataAggregationHelper.ts b/src/abis/umbrellaDataAggregationHelper.ts
index 90b4419..5928841 100644
--- a/src/abis/umbrellaDataAggregationHelper.ts
+++ b/src/abis/umbrellaDataAggregationHelper.ts
@@ -1,12 +1,88 @@
 export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
+  {
+    inputs: [
+      { internalType: "address", name: "rewardsController_", type: "address" },
+      { internalType: "address", name: "owner_", type: "address" },
+    ],
+    stateMutability: "nonpayable",
+    type: "constructor",
+  },
+  { inputs: [], name: "EthTransferFailed", type: "error" },
+  { inputs: [], name: "OnlyRescueGuardian", type: "error" },
+  { inputs: [{ internalType: "address", name: "owner", type: "address" }], name: "OwnableInvalidOwner", type: "error" },
+  {
+    inputs: [{ internalType: "address", name: "account", type: "address" }],
+    name: "OwnableUnauthorizedAccount",
+    type: "error",
+  },
+  {
+    inputs: [{ internalType: "address", name: "token", type: "address" }],
+    name: "SafeERC20FailedOperation",
+    type: "error",
+  },
+  { inputs: [], name: "ZeroAddress", type: "error" },
+  {
+    anonymous: false,
+    inputs: [
+      { indexed: true, internalType: "address", name: "caller", type: "address" },
+      { indexed: true, internalType: "address", name: "token", type: "address" },
+      { indexed: true, internalType: "address", name: "to", type: "address" },
+      { indexed: false, internalType: "uint256", name: "amount", type: "uint256" },
+    ],
+    name: "ERC20Rescued",
+    type: "event",
+  },
+  {
+    anonymous: false,
+    inputs: [
+      { indexed: true, internalType: "address", name: "caller", type: "address" },
+      { indexed: true, internalType: "address", name: "to", type: "address" },
+      { indexed: false, internalType: "uint256", name: "amount", type: "uint256" },
+    ],
+    name: "NativeTokensRescued",
+    type: "event",
+  },
+  {
+    anonymous: false,
+    inputs: [
+      { indexed: true, internalType: "address", name: "previousOwner", type: "address" },
+      { indexed: true, internalType: "address", name: "newOwner", type: "address" },
+    ],
+    name: "OwnershipTransferred",
+    type: "event",
+  },
+  {
+    inputs: [],
+    name: "REWARDS_CONTROLLER",
+    outputs: [{ internalType: "contract IRewardsController", name: "", type: "address" }],
+    stateMutability: "view",
+    type: "function",
+  },
+  {
+    inputs: [
+      { internalType: "address", name: "to", type: "address" },
+      { internalType: "uint256", name: "amount", type: "uint256" },
+    ],
+    name: "emergencyEtherTransfer",
+    outputs: [],
+    stateMutability: "nonpayable",
+    type: "function",
+  },
+  {
+    inputs: [
+      { internalType: "address", name: "erc20Token", type: "address" },
+      { internalType: "address", name: "to", type: "address" },
+      { internalType: "uint256", name: "amount", type: "uint256" },
+    ],
+    name: "emergencyTokenTransfer",
+    outputs: [],
+    stateMutability: "nonpayable",
+    type: "function",
+  },
   {
     inputs: [
       { internalType: "contract IUmbrella", name: "umbrella", type: "address" },
-      {
-        internalType: "contract IAaveOracle",
-        name: "aaveOracle",
-        type: "address",
-      },
+      { internalType: "contract IAaveOracle", name: "aaveOracle", type: "address" },
       { internalType: "address", name: "user", type: "address" },
     ],
     name: "getAllAggregatedData",
@@ -26,6 +102,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
             type: "tuple",
           },
           { internalType: "uint256", name: "totalAssets", type: "uint256" },
+          { internalType: "uint256", name: "maxTotalAssets", type: "uint256" },
           { internalType: "uint256", name: "targetLiquidity", type: "uint256" },
           { internalType: "address", name: "reserve", type: "address" },
           { internalType: "bool", name: "isStakeConfigured", type: "bool" },
@@ -43,16 +120,8 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
                 name: "rewardTokenData",
                 type: "tuple",
               },
-              {
-                internalType: "uint256",
-                name: "maxEmissionPerSecond",
-                type: "uint256",
-              },
-              {
-                internalType: "uint256",
-                name: "distributionEnd",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "maxEmissionPerSecond", type: "uint256" },
+              { internalType: "uint256", name: "distributionEnd", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.RewardTokenData[]",
             name: "rewardsTokenData",
@@ -68,11 +137,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           { internalType: "address", name: "stakeToken", type: "address" },
           {
             components: [
-              {
-                internalType: "enum DataAggregationHelper.TokenType",
-                name: "typeOfToken",
-                type: "uint8",
-              },
+              { internalType: "enum DataAggregationHelper.TokenType", name: "typeOfToken", type: "uint8" },
               {
                 components: [
                   { internalType: "address", name: "token", type: "address" },
@@ -98,19 +163,11 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
       {
         components: [
           { internalType: "address", name: "stakeToken", type: "address" },
-          {
-            internalType: "uint256",
-            name: "stakeUserBalance",
-            type: "uint256",
-          },
+          { internalType: "uint256", name: "stakeUserBalance", type: "uint256" },
           {
             components: [
               { internalType: "address", name: "reward", type: "address" },
-              {
-                internalType: "uint256",
-                name: "currentReward",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "currentReward", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.RewardTokenUserData[]",
             name: "rewardsTokenUserData",
@@ -126,11 +183,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           { internalType: "address", name: "stakeToken", type: "address" },
           {
             components: [
-              {
-                internalType: "enum DataAggregationHelper.TokenType",
-                name: "typeOfToken",
-                type: "uint8",
-              },
+              { internalType: "enum DataAggregationHelper.TokenType", name: "typeOfToken", type: "uint8" },
               { internalType: "address", name: "token", type: "address" },
               { internalType: "uint256", name: "userBalance", type: "uint256" },
             ],
@@ -150,11 +203,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
   {
     inputs: [
       { internalType: "contract IUmbrella", name: "umbrella", type: "address" },
-      {
-        internalType: "contract IAaveOracle",
-        name: "aaveOracle",
-        type: "address",
-      },
+      { internalType: "contract IAaveOracle", name: "aaveOracle", type: "address" },
     ],
     name: "getTokensAggregatedData",
     outputs: [
@@ -173,6 +222,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
             type: "tuple",
           },
           { internalType: "uint256", name: "totalAssets", type: "uint256" },
+          { internalType: "uint256", name: "maxTotalAssets", type: "uint256" },
           { internalType: "uint256", name: "targetLiquidity", type: "uint256" },
           { internalType: "address", name: "reserve", type: "address" },
           { internalType: "bool", name: "isStakeConfigured", type: "bool" },
@@ -190,16 +240,8 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
                 name: "rewardTokenData",
                 type: "tuple",
               },
-              {
-                internalType: "uint256",
-                name: "maxEmissionPerSecond",
-                type: "uint256",
-              },
-              {
-                internalType: "uint256",
-                name: "distributionEnd",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "maxEmissionPerSecond", type: "uint256" },
+              { internalType: "uint256", name: "distributionEnd", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.RewardTokenData[]",
             name: "rewardsTokenData",
@@ -217,11 +259,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
   {
     inputs: [
       { internalType: "contract IUmbrella", name: "umbrella", type: "address" },
-      {
-        internalType: "contract IAaveOracle",
-        name: "aaveOracle",
-        type: "address",
-      },
+      { internalType: "contract IAaveOracle", name: "aaveOracle", type: "address" },
     ],
     name: "getTokensRouteData",
     outputs: [
@@ -230,11 +268,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           { internalType: "address", name: "stakeToken", type: "address" },
           {
             components: [
-              {
-                internalType: "enum DataAggregationHelper.TokenType",
-                name: "typeOfToken",
-                type: "uint8",
-              },
+              { internalType: "enum DataAggregationHelper.TokenType", name: "typeOfToken", type: "uint8" },
               {
                 components: [
                   { internalType: "address", name: "token", type: "address" },
@@ -264,11 +298,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
   {
     inputs: [
       { internalType: "contract IUmbrella", name: "umbrella", type: "address" },
-      {
-        internalType: "contract IAaveOracle",
-        name: "aaveOracle",
-        type: "address",
-      },
+      { internalType: "contract IAaveOracle", name: "aaveOracle", type: "address" },
     ],
     name: "getUmbrellaData",
     outputs: [
@@ -289,13 +319,10 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
                 type: "tuple",
               },
               { internalType: "uint256", name: "totalAssets", type: "uint256" },
+              { internalType: "uint256", name: "maxTotalAssets", type: "uint256" },
               { internalType: "uint256", name: "totalSupply", type: "uint256" },
               { internalType: "uint256", name: "cooldown", type: "uint256" },
-              {
-                internalType: "uint256",
-                name: "unstakeWindow",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "unstakeWindow", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.StakeTokenConfig",
             name: "stakeTokenConfig",
@@ -303,51 +330,23 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           },
           {
             components: [
-              {
-                internalType: "uint256",
-                name: "targetLiquidity",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "targetLiquidity", type: "uint256" },
               {
                 components: [
                   {
                     components: [
-                      {
-                        internalType: "address",
-                        name: "token",
-                        type: "address",
-                      },
-                      {
-                        internalType: "uint256",
-                        name: "price",
-                        type: "uint256",
-                      },
+                      { internalType: "address", name: "token", type: "address" },
+                      { internalType: "uint256", name: "price", type: "uint256" },
                       { internalType: "string", name: "name", type: "string" },
-                      {
-                        internalType: "string",
-                        name: "symbol",
-                        type: "string",
-                      },
-                      {
-                        internalType: "uint8",
-                        name: "decimals",
-                        type: "uint8",
-                      },
+                      { internalType: "string", name: "symbol", type: "string" },
+                      { internalType: "uint8", name: "decimals", type: "uint8" },
                     ],
                     internalType: "struct DataAggregationHelper.TokenData",
                     name: "reward",
                     type: "tuple",
                   },
-                  {
-                    internalType: "uint256",
-                    name: "maxEmissionPerSecond",
-                    type: "uint256",
-                  },
-                  {
-                    internalType: "uint256",
-                    name: "distributionEnd",
-                    type: "uint256",
-                  },
+                  { internalType: "uint256", name: "maxEmissionPerSecond", type: "uint256" },
+                  { internalType: "uint256", name: "distributionEnd", type: "uint256" },
                 ],
                 internalType: "struct DataAggregationHelper.RewardConfig[]",
                 name: "rewardConfigs",
@@ -361,16 +360,8 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           {
             components: [
               { internalType: "address", name: "reserve", type: "address" },
-              {
-                internalType: "uint256",
-                name: "deficitOffset",
-                type: "uint256",
-              },
-              {
-                internalType: "uint256",
-                name: "pendingDeficit",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "deficitOffset", type: "uint256" },
+              { internalType: "uint256", name: "pendingDeficit", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.UmbrellaConfig",
             name: "umbrellaConfig",
@@ -395,19 +386,11 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
       {
         components: [
           { internalType: "address", name: "stakeToken", type: "address" },
-          {
-            internalType: "uint256",
-            name: "stakeUserBalance",
-            type: "uint256",
-          },
+          { internalType: "uint256", name: "stakeUserBalance", type: "uint256" },
           {
             components: [
               { internalType: "address", name: "reward", type: "address" },
-              {
-                internalType: "uint256",
-                name: "currentReward",
-                type: "uint256",
-              },
+              { internalType: "uint256", name: "currentReward", type: "uint256" },
             ],
             internalType: "struct DataAggregationHelper.RewardTokenUserData[]",
             name: "rewardsTokenUserData",
@@ -434,11 +417,7 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
           { internalType: "address", name: "stakeToken", type: "address" },
           {
             components: [
-              {
-                internalType: "enum DataAggregationHelper.TokenType",
-                name: "typeOfToken",
-                type: "uint8",
-              },
+              { internalType: "enum DataAggregationHelper.TokenType", name: "typeOfToken", type: "uint8" },
               { internalType: "address", name: "token", type: "address" },
               { internalType: "uint256", name: "userBalance", type: "uint256" },
             ],
@@ -455,4 +434,33 @@ export const UMBRELLA_DATA_AGGREGATION_HELPER_ABI = [
     stateMutability: "view",
     type: "function",
   },
+  {
+    inputs: [{ internalType: "address", name: "", type: "address" }],
+    name: "maxRescue",
+    outputs: [{ internalType: "uint256", name: "", type: "uint256" }],
+    stateMutability: "pure",
+    type: "function",
+  },
+  {
+    inputs: [],
+    name: "owner",
+    outputs: [{ internalType: "address", name: "", type: "address" }],
+    stateMutability: "view",
+    type: "function",
+  },
+  { inputs: [], name: "renounceOwnership", outputs: [], stateMutability: "nonpayable", type: "function" },
+  {
+    inputs: [{ internalType: "address", name: "newOwner", type: "address" }],
+    name: "transferOwnership",
+    outputs: [],
+    stateMutability: "nonpayable",
+    type: "function",
+  },
+  {
+    inputs: [],
+    name: "whoCanRescue",
+    outputs: [{ internalType: "address", name: "", type: "address" }],
+    stateMutability: "view",
+    type: "function",
+  },
 ] as const;
diff --git a/src/constants/markets.ts b/src/constants/markets.ts
index b3e27b7..ab32c07 100644
--- a/src/constants/markets.ts
+++ b/src/constants/markets.ts
@@ -28,8 +28,10 @@ export const MARKETS: Market[] = [
           oracle: AaveV3BaseSepolia.ORACLE,
           rewardsController: UmbrellaBaseSepolia.UMBRELLA_REWARDS_CONTROLLER,
           umbrellaHelper: UmbrellaBaseSepolia.UMBRELLA,
-          batchHelper: UmbrellaBaseSepolia.UMBRELLA_BATCH_HELPER,
-          umbrellaDataAggregationHelper: UmbrellaBaseSepolia.DATA_AGGREGATION_HELPER,
+          // batchHelper: UmbrellaBaseSepolia.UMBRELLA_BATCH_HELPER,
+          batchHelper: "0x32eb722eF6fC2B5ea50e833Fb538769fF09A6C1c", // New version
+          // umbrellaDataAggregationHelper: UmbrellaBaseSepolia.DATA_AGGREGATION_HELPER,
+          umbrellaDataAggregationHelper: "0xDdE9091B6F9003f3b7070Ce7c377Eee09a8e0245", // New version
           wrapNativeTokenAddress: AaveV3BaseSepolia.ASSETS.WETH.UNDERLYING,
         } as const,
       ]

From e1257e37867817118fc60d66c7746c50582cc99e Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 09:06:51 +0100
Subject: [PATCH 2/6] feat: add maxTargetLiquidity

---
 src/hooks/useAllStkTokens.ts | 3 ++-
 src/types/token.ts           | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/hooks/useAllStkTokens.ts b/src/hooks/useAllStkTokens.ts
index 9011a50..3b842de 100644
--- a/src/hooks/useAllStkTokens.ts
+++ b/src/hooks/useAllStkTokens.ts
@@ -53,7 +53,7 @@ export const useAllStkTokens = () => {
       const [aggregatedData, pathData, userAggregatedData, userPathData] = data;
 
       return aggregatedData
-        .map(({ stakeTokenData, rewardsTokenData, totalAssets, targetLiquidity }, index) => {
+        .map(({ stakeTokenData, rewardsTokenData, totalAssets, targetLiquidity, maxTotalAssets }, index) => {
           const { stakeUserBalance, rewardsTokenUserData } = userAggregatedData[index] ?? {};
           const stkTokenPathData = pathData[index] as StkTokenPathData;
           const stkTokenUserPathData = userPathData[index] as UserPathData;
@@ -121,6 +121,7 @@ export const useAllStkTokens = () => {
             totalAssets,
             totalRewardsUSDAmount: rewards.reduce((acc, { usdAmount }) => acc + (usdAmount ?? 0), 0),
             targetLiquidity,
+            maxTotalAssets,
             rewards,
             reserve: reserve ?? null,
             stata:
diff --git a/src/types/token.ts b/src/types/token.ts
index 72ee971..7367412 100644
--- a/src/types/token.ts
+++ b/src/types/token.ts
@@ -95,6 +95,7 @@ export type StkToken = Token &
     totalAssets: bigint;
     totalRewardsUSDAmount: number;
     targetLiquidity: bigint;
+    maxTotalAssets: bigint;
     isUnderlyingStataToken: boolean;
     rewards: Reward[];
     reserve: Reserve | null;

From 9c09a6a17eb466c04a540d194a53acec519c911f Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 13:23:57 +0100
Subject: [PATCH 3/6] feat: add to supply cap to summary

---
 .../StakedSummary/StakedSummary.tsx           | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/src/app/components/StakedSummary/StakedSummary.tsx b/src/app/components/StakedSummary/StakedSummary.tsx
index b546dfe..93d4be0 100644
--- a/src/app/components/StakedSummary/StakedSummary.tsx
+++ b/src/app/components/StakedSummary/StakedSummary.tsx
@@ -42,6 +42,7 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
       decimals,
       latestAnswerFormatted,
       targetLiquidity,
+      maxTotalAssets,
       underlying,
       isUnderlyingStataToken,
       apyData,
@@ -49,6 +50,7 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
 
     const stakedUSD = Number(formatUnits(totalAssets, decimals)) * latestAnswerFormatted;
     const targetLiquidityUSD = Number(formatUnits(targetLiquidity, decimals)) * latestAnswerFormatted;
+    const maxTotalAssetsUSD = Number(formatUnits(maxTotalAssets, decimals)) * latestAnswerFormatted;
 
     return (
       <div key={address} className="border-main-200 flex flex-col gap-4 border-b pb-4 last:border-b-0 last:pb-0">
@@ -71,6 +73,10 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
             <div className="text-main-500">Target Liquidity</div>
             <NumberDisplay value={targetLiquidityUSD} type="currency" className="font-bold" />
           </div>
+          <div className="flex flex-col gap-1">
+            <div className="text-main-500">Supply Cap</div>
+            <NumberDisplay value={maxTotalAssetsUSD} type="currency" className="font-bold" />
+          </div>
           <div className="flex flex-col gap-1">
             <div className="text-main-500">Current Yield</div>
             <NumberDisplay value={apyData.total} type="percent" className="font-bold" />
@@ -102,6 +108,7 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
       underlying,
       isUnderlyingStataToken,
       apyData,
+      maxTotalAssets,
     } = selectedToken;
 
     const rewardColors = ["bg-green-300", "bg-orange-300", "bg-red-300", "bg-pink-300", "bg-yellow-300"];
@@ -110,6 +117,7 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
     const targetLiquidityUSD = Number(formatUnits(targetLiquidity, decimals)) * latestAnswerFormatted;
 
     const currentUmbrellaData = umbrellaData?.find((data) => data.stakeTokenConfig.stakeToken.address === address);
+    const maxTotalAssetsUSD = Number(formatUnits(maxTotalAssets, decimals)) * latestAnswerFormatted;
 
     return (
       <ModalRoot open={!!selectedToken} onOpenChange={(open) => !open && setSelectedToken(null)}>
@@ -134,6 +142,10 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
                     <div className="text-main-500 text-sm font-medium">Total Staked</div>
                     <NumberDisplay value={stakedUSD} type="currency" className="text-lg font-bold" />
                   </div>
+                  <div className="flex flex-col gap-2">
+                    <div className="text-main-500 text-sm font-medium">Supply Cap</div>
+                    <NumberDisplay value={maxTotalAssetsUSD} type="currency" className="text-lg font-bold" />
+                  </div>
 
                   <div className="flex flex-col gap-4">
                     <div className="flex flex-col gap-2">
@@ -342,12 +354,14 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
                   decimals,
                   latestAnswerFormatted,
                   targetLiquidity,
+                  maxTotalAssets,
                   underlying,
                   isUnderlyingStataToken,
                   apyData,
                 }) => {
                   const stakedUSD = Number(formatUnits(totalAssets, decimals)) * latestAnswerFormatted;
                   const targetLiquidityUSD = Number(formatUnits(targetLiquidity, decimals)) * latestAnswerFormatted;
+                  const maxTotalAssetsUSD = Number(formatUnits(maxTotalAssets, decimals)) * latestAnswerFormatted;
 
                   return (
                     <Tooltip
@@ -395,6 +409,12 @@ export const StakedSummary = memo(({ stkTokens }: StakedSummaryProps) => {
                               <NumberDisplay value={targetLiquidityUSD} type="currency" className="font-bold" />
                             </div>
                           </div>
+                          <div className="flex items-center justify-between gap-2">
+                            <div className="dark:text-main-500">Supply Cap</div>
+                            <div className="text-main-900 dark:text-white">
+                              <NumberDisplay value={maxTotalAssetsUSD} type="currency" className="font-bold" />
+                            </div>
+                          </div>
                           <div className="bg-main-100 dark:bg-main-600 my-2 h-px w-full" />
                           <div className="flex items-center justify-between gap-2">
                             <div className="dark:text-main-500">Current Yield</div>

From 28cb1e0227461ef345e2a5d1cc44c520ff93844c Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 13:30:22 +0100
Subject: [PATCH 4/6] feat: edgecase for extra-large numbers

---
 src/components/NumberDisplay/NumberDisplay.tsx | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/components/NumberDisplay/NumberDisplay.tsx b/src/components/NumberDisplay/NumberDisplay.tsx
index bfabee9..dba0809 100644
--- a/src/components/NumberDisplay/NumberDisplay.tsx
+++ b/src/components/NumberDisplay/NumberDisplay.tsx
@@ -25,6 +25,11 @@ export const NumberDisplay = ({
       prefix = customPrefix ? `${customPrefix}<` : "<";
     }
 
+    if (value > 1e15) {
+      displayValue = 1e15;
+      prefix = customPrefix ? `${customPrefix}>` : ">";
+    }
+
     return (
       <NumberFlow
         value={displayValue}
@@ -51,6 +56,11 @@ export const NumberDisplay = ({
       prefix = customPrefix ? `${customPrefix}<` : "<";
     }
 
+    if (value > 1e15) {
+      displayValue = 1e15;
+      prefix = customPrefix ? `${customPrefix}>` : ">";
+    }
+
     return (
       <NumberFlow
         value={displayValue / 100}

From 0a18d2828ca98d5f3b1950c6d707387ad68e0b46 Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 16:10:34 +0100
Subject: [PATCH 5/6] feat: support immediate withdrawals

---
 src/app/components/UmbrellaTable/Actions.tsx             | 2 +-
 src/app/withdraw/[address]/page.tsx                      | 2 +-
 .../useAllUmbrellaCooldowns/useUmbrellaCooldownStatus.ts | 9 ++++++++-
 src/types/cooldown.ts                                    | 2 +-
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/app/components/UmbrellaTable/Actions.tsx b/src/app/components/UmbrellaTable/Actions.tsx
index d048d66..4a4586e 100644
--- a/src/app/components/UmbrellaTable/Actions.tsx
+++ b/src/app/components/UmbrellaTable/Actions.tsx
@@ -74,7 +74,7 @@ export const Actions = ({ token }: ActionsProps) => {
 
   return (
     <div className="flex flex-col gap-3 md:flex-row md:items-center md:justify-center md:gap-4">
-      {status === "withdraw" && (
+      {(status === "withdraw" || status === "immediate-withdraw") && (
         <Button href={`/withdraw/${address}`} prefetch primary elevation={1} size="lg" className="gap-1 font-semibold">
           <ChevronsDownIcon size={16} />
           Withdraw
diff --git a/src/app/withdraw/[address]/page.tsx b/src/app/withdraw/[address]/page.tsx
index b467cec..fc2df67 100644
--- a/src/app/withdraw/[address]/page.tsx
+++ b/src/app/withdraw/[address]/page.tsx
@@ -26,7 +26,7 @@ export default function WithdrawPage(props: WithdrawPageProps) {
   const stkToken = stkTokens?.find((token) => token.address === address);
 
   useEffect(() => {
-    if (status !== "withdraw") {
+    if (status !== "withdraw" && status !== "immediate-withdraw") {
       router.push("/");
     }
   }, [router, status]);
diff --git a/src/hooks/useAllUmbrellaCooldowns/useUmbrellaCooldownStatus.ts b/src/hooks/useAllUmbrellaCooldowns/useUmbrellaCooldownStatus.ts
index e9a472b..9109568 100644
--- a/src/hooks/useAllUmbrellaCooldowns/useUmbrellaCooldownStatus.ts
+++ b/src/hooks/useAllUmbrellaCooldowns/useUmbrellaCooldownStatus.ts
@@ -1,4 +1,5 @@
 import { useUmbrellaCooldownData } from "@/hooks/useAllUmbrellaCooldowns/useUmbrellaCooldownData";
+import { useUmbrellaData } from "@/hooks/useUmbrellaData";
 import { Address } from "viem";
 import { useEffect, useState } from "react";
 import { CooldownStatus } from "@/types/cooldown";
@@ -25,8 +26,14 @@ export const useUmbrellaCooldownStatus = (address: Address) => {
   const setCooldownStatus = useDashboardStore.use.setCooldownStatus();
 
   const { data } = useUmbrellaCooldownData(address);
+  const { data: umbrellaData } = useUmbrellaData();
   const { endOfCooldown, withdrawalEndsAt } = data || {};
-  const status = calculateStatus(now, endOfCooldown, withdrawalEndsAt);
+
+  const configuredCooldownPeriod = umbrellaData?.find((item) => item.stakeTokenConfig.stakeToken.address === address)
+    ?.stakeTokenConfig.cooldown;
+
+  const status =
+    configuredCooldownPeriod === 0 ? "immediate-withdraw" : calculateStatus(now, endOfCooldown, withdrawalEndsAt);
 
   useEffect(() => {
     const updateStatus = () => {
diff --git a/src/types/cooldown.ts b/src/types/cooldown.ts
index d928a9c..b9086eb 100644
--- a/src/types/cooldown.ts
+++ b/src/types/cooldown.ts
@@ -1 +1 @@
-export type CooldownStatus = "cooldown" | "withdraw" | "none";
+export type CooldownStatus = "cooldown" | "withdraw" | "none" | "immediate-withdraw";

From 47c49a5ebee695636f432e35a71f01735898c1a2 Mon Sep 17 00:00:00 2001
From: "K.K." <m@mnm.mn>
Date: Mon, 1 Sep 2025 18:01:48 +0100
Subject: [PATCH 6/6] feat: account for maxTotalAssets in calculations

---
 .../StakeATokenForm/StakeATokenForm.tsx       |  3 +-
 .../StakeStataTokenForm.tsx                   |  3 +-
 .../StakeUnderlyingTokenForm.tsx              |  4 +-
 src/utils/calculations/index.ts               |  4 +-
 src/utils/calculations/reserves.ts            | 48 +++++++++++++++++--
 src/utils/calculations/staking.ts             | 10 ++++
 6 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/src/app/stake/a/[address]/components/StakeATokenForm/StakeATokenForm.tsx b/src/app/stake/a/[address]/components/StakeATokenForm/StakeATokenForm.tsx
index eb1e4c4..878e1b9 100644
--- a/src/app/stake/a/[address]/components/StakeATokenForm/StakeATokenForm.tsx
+++ b/src/app/stake/a/[address]/components/StakeATokenForm/StakeATokenForm.tsx
@@ -14,6 +14,7 @@ import { useSafeApproveAndStake } from "@/hooks/useSafeApproveAndStake";
 import { useStake } from "@/hooks/useStake";
 import { useTxFormSignature } from "@/providers/TxFormProvider/TxFormContext";
 import { Reserve, StkToken } from "@/types/token";
+import { calculateMaxStakeAmount } from "@/utils/calculations";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { useQueryClient } from "@tanstack/react-query";
 import { useMemo } from "react";
@@ -37,7 +38,7 @@ export const StakeATokenForm = ({ asset, stkToken, reserves }: StakeATokenFormPr
 
   const reserve = stkToken.reserve!;
   const { reserveId, balance, decimals, symbol, latestAnswer } = reserve!;
-  const maxAmount = balance || 0n;
+  const maxAmount = calculateMaxStakeAmount(stkToken, balance || 0n);
 
   const schema = useMemo(
     () =>
diff --git a/src/app/stake/stata/[address]/components/StakeStataTokenForm/StakeStataTokenForm.tsx b/src/app/stake/stata/[address]/components/StakeStataTokenForm/StakeStataTokenForm.tsx
index 9ea2356..eb95ef8 100644
--- a/src/app/stake/stata/[address]/components/StakeStataTokenForm/StakeStataTokenForm.tsx
+++ b/src/app/stake/stata/[address]/components/StakeStataTokenForm/StakeStataTokenForm.tsx
@@ -15,6 +15,7 @@ import { useSafeApproveAndStake } from "@/hooks/useSafeApproveAndStake";
 import { useStake } from "@/hooks/useStake";
 import { useTxFormSignature } from "@/providers/TxFormProvider/TxFormContext";
 import { StataToken, StkToken } from "@/types/token";
+import { calculateMaxStakeAmount } from "@/utils/calculations";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { useMemo } from "react";
 import { Controller, FormProvider, useForm } from "react-hook-form";
@@ -39,7 +40,7 @@ export const StakeStataTokenForm = ({ asset, stkToken, stataToken }: StakeStataT
     error: approveAndStakeError,
   } = useSafeApproveAndStake();
 
-  const maxAmount = stataToken.balance || 0n;
+  const maxAmount = calculateMaxStakeAmount(stkToken, stataToken.balance || 0n);
   const { decimals, latestAnswer } = stataToken;
 
   const schema = useMemo(() => createStakeStataFormSchema({ maxAmount, isSafeWallet }), [maxAmount, isSafeWallet]);
diff --git a/src/app/stake/underlying/[address]/components/StakeUnderlyingTokenForm/StakeUnderlyingTokenForm.tsx b/src/app/stake/underlying/[address]/components/StakeUnderlyingTokenForm/StakeUnderlyingTokenForm.tsx
index ff0690a..ce79866 100644
--- a/src/app/stake/underlying/[address]/components/StakeUnderlyingTokenForm/StakeUnderlyingTokenForm.tsx
+++ b/src/app/stake/underlying/[address]/components/StakeUnderlyingTokenForm/StakeUnderlyingTokenForm.tsx
@@ -15,7 +15,7 @@ import { useSafeApproveAndStake } from "@/hooks/useSafeApproveAndStake";
 import { useStake } from "@/hooks/useStake";
 import { useTxFormSignature } from "@/providers/TxFormProvider/TxFormContext";
 import { StkToken } from "@/types/token";
-import { calculateMaxSupply } from "@/utils/calculations";
+import { calculateMaxUnderlyingStake } from "@/utils/calculations";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { useMemo } from "react";
 import { Controller, FormProvider, useForm } from "react-hook-form";
@@ -42,7 +42,7 @@ export const StakeUnderlyingTokenForm = ({ asset, stkToken, totalSupplied }: Sta
   const { stake, data: hash, isPending: isTxPending, error: depositError } = useStake();
 
   const { decimals, latestAnswer, symbol } = stkToken.underlying;
-  const maxAmount = calculateMaxSupply(stkToken, totalSupplied);
+  const maxAmount = calculateMaxUnderlyingStake(stkToken, totalSupplied);
 
   const schema = useMemo(() => createStakeFormSchema({ maxAmount, isSafeWallet }), [maxAmount, isSafeWallet]);
   const formMethods = useForm<StakeUnderlyingFormValues>({
diff --git a/src/utils/calculations/index.ts b/src/utils/calculations/index.ts
index e52f80c..c97e2f4 100644
--- a/src/utils/calculations/index.ts
+++ b/src/utils/calculations/index.ts
@@ -4,6 +4,6 @@ export { calculateHealthFactor, calculateNewHealthFactor, getHealthStatus } from
 export type { CalculateHealthFactorArgs, CalculateNewHealthFactorArgs, HealthFactorStatus } from "./healthFactor";
 export { calculateAPYEarnings, calculateSimpleInterest } from "./interest";
 export { calculateMetrics } from "./metrics";
-export { calculateMaxSupply } from "./reserves";
+export { calculateMaxSupply, calculateMaxUnderlyingStake } from "./reserves";
 export { sumUpAllRewards } from "./rewards";
-export { calculateAvailableToStakeUsd } from "./staking";
+export { calculateAvailableToStakeUsd, calculateMaxStakeAmount } from "./staking";
diff --git a/src/utils/calculations/reserves.ts b/src/utils/calculations/reserves.ts
index 34dc0f6..f529710 100644
--- a/src/utils/calculations/reserves.ts
+++ b/src/utils/calculations/reserves.ts
@@ -2,9 +2,13 @@ import { parseUnits } from "viem";
 import { StkToken } from "@/types/token";
 
 export const calculateMaxSupply = (stkToken: StkToken, totalSupplied?: bigint) => {
-  const { reserve, underlying } = stkToken;
+  const { reserve, underlying, totalAssets, maxTotalAssets } = stkToken;
   const { balance = 0n } = underlying;
 
+  if (totalAssets >= maxTotalAssets) {
+    return 0n;
+  }
+
   if (!reserve) {
     return balance;
   }
@@ -17,8 +21,44 @@ export const calculateMaxSupply = (stkToken: StkToken, totalSupplied?: bigint) =
 
   const scaledSupplyCap = parseUnits(String(supplyCap), decimals);
   const totalSupplyWithAccrued = totalSupplied + accruedToTreasury;
-  const availableSupply =
-    scaledSupplyCap > totalSupplyWithAccrued ? scaledSupplyCap - totalSupplyWithAccrued : 0n;
+  const availableSupply = scaledSupplyCap > totalSupplyWithAccrued ? scaledSupplyCap - totalSupplyWithAccrued : 0n;
+
+  const availableUmbrellaSupply = maxTotalAssets - totalAssets;
+
+  const maxAvailableSupply = availableSupply >= balance ? balance : availableSupply;
+  const maxAvailableUmbrellaSupply = availableUmbrellaSupply >= balance ? balance : availableUmbrellaSupply;
+
+  const result = maxAvailableSupply < maxAvailableUmbrellaSupply ? maxAvailableSupply : maxAvailableUmbrellaSupply;
+
+  return result;
+};
+
+export const calculateMaxUnderlyingStake = (stkToken: StkToken, totalSupplied?: bigint) => {
+  const { reserve, underlying, totalAssets, maxTotalAssets } = stkToken;
+  const { balance = 0n } = underlying;
+
+  if (!reserve) {
+    return balance;
+  }
+
+  const { decimals, supplyCap, accruedToTreasury } = reserve;
+
+  let availableAaveSupply = balance;
+  if (supplyCap > 0n && totalSupplied) {
+    const scaledSupplyCap = parseUnits(String(supplyCap), decimals);
+    const totalSupplyWithAccrued = totalSupplied + accruedToTreasury;
+    const aaveAvailable = scaledSupplyCap > totalSupplyWithAccrued ? scaledSupplyCap - totalSupplyWithAccrued : 0n;
+    availableAaveSupply = aaveAvailable >= balance ? balance : aaveAvailable;
+  }
+
+  let availableUmbrellaSupply = balance;
+  if (totalAssets < maxTotalAssets) {
+    const umbrellaAvailable = maxTotalAssets - totalAssets;
+    availableUmbrellaSupply = umbrellaAvailable >= balance ? balance : umbrellaAvailable;
+  } else {
+    availableUmbrellaSupply = 0n;
+  }
 
-  return availableSupply >= balance ? balance : availableSupply;
+  const result = availableAaveSupply < availableUmbrellaSupply ? availableAaveSupply : availableUmbrellaSupply;
+  return result;
 };
diff --git a/src/utils/calculations/staking.ts b/src/utils/calculations/staking.ts
index 384d927..7fa131e 100644
--- a/src/utils/calculations/staking.ts
+++ b/src/utils/calculations/staking.ts
@@ -1,4 +1,5 @@
 import { Asset } from "@/types/token";
+import { StkToken } from "@/types/token";
 
 export const calculateAvailableToStakeUsd = (assets: Asset[]) => {
   return assets.reduce((acc, asset) => {
@@ -9,3 +10,12 @@ export const calculateAvailableToStakeUsd = (assets: Asset[]) => {
     return acc;
   }, 0);
 };
+
+export const calculateMaxStakeAmount = (stkToken: StkToken, userBalance: bigint) => {
+  const { totalAssets, maxTotalAssets } = stkToken;
+
+  const availableUmbrellaSupply = maxTotalAssets > totalAssets ? maxTotalAssets - totalAssets : 0n;
+  const result = availableUmbrellaSupply >= userBalance ? userBalance : availableUmbrellaSupply;
+
+  return result;
+};