Fix balance related blockers for migration

frontend/components/MainPage/index.tsx

@@ -19,25 +19,6 @@ export const Main = () => {
   const isStakingContractSectionEnabled = useFeatureFlag(
     'staking-contract-section',
   );
-  // const { refetch: updateServicesState } = useServices();
-  // const {
-  //   updateBalances,
-  //   isLoaded: isBalanceLoaded,
-  //   setIsLoaded: setIsBalanceLoaded,
-  // } = useBalanceContext();
-
-  // TODO: reintroduce later,  non critical
-  // useEffect(() => {
-  //   if (!isBalanceLoaded) {
-  //     updateServicesState?.().then(() => updateBalances());
-  //     setIsBalanceLoaded(true);
-  //   }
-  // }, [
-  //   isBalanceLoaded,
-  //   setIsBalanceLoaded,
-  //   updateBalances,
-  //   updateServicesState,
-  // ]);
 
   return (
     <Card

frontend/components/MainPage/sections/OlasBalanceSection.tsx

@@ -26,7 +26,7 @@ const Balance = styled.span`
 `;
 
 export const MainOlasBalance = () => {
-  const { selectedService } = useServices();
+  const { selectedService, selectedAgentConfig } = useServices();
   const { isLoaded: isBalanceLoaded } = useBalanceContext();
   const { masterWalletBalances } = useMasterBalances();
   const { serviceStakedBalances, serviceWalletBalances } = useServiceBalances(
@@ -36,31 +36,40 @@ export const MainOlasBalance = () => {
   const isBalanceBreakdownEnabled = useFeatureFlag('manage-wallet');
 
   const displayedBalance = useMemo(() => {
-    // olas across master wallets, safes and eoa
+    // olas across master wallet -- safes and eoa -- on relevant chains for agent
     const masterWalletOlasBalance = masterWalletBalances?.reduce(
-      (acc, { symbol, balance }) => {
-        if (symbol === TokenSymbol.OLAS) {
+      (acc, { symbol, balance, evmChainId }) => {
+        if (
+          symbol === TokenSymbol.OLAS &&
+          selectedAgentConfig.requiresAgentSafesOn.includes(evmChainId)
+        ) {
           return acc + Number(balance);
         }
         return acc;
       },
       0,
     );
 
-    // olas across all service wallets
+    // olas across all wallets owned by selected service
     const serviceWalletOlasBalance = serviceWalletBalances?.reduce(
-      (acc, { symbol, balance }) => {
-        if (symbol === TokenSymbol.OLAS) {
+      (acc, { symbol, balance, evmChainId }) => {
+        if (
+          symbol === TokenSymbol.OLAS &&
+          selectedAgentConfig.requiresAgentSafesOn.includes(evmChainId)
+        ) {
           return acc + Number(balance);
         }
         return acc;
       },
       0,
     );
 
-    // olas staked across all services
+    // olas staked across services on relevant chains for agent
     const serviceStakedOlasBalance = serviceStakedBalances?.reduce(
-      (acc, { olasBondBalance, olasDepositBalance }) => {
+      (acc, { olasBondBalance, olasDepositBalance, evmChainId }) => {
+        if (!selectedAgentConfig.requiresAgentSafesOn.includes(evmChainId)) {
+          return acc;
+        }
         return acc + Number(olasBondBalance) + Number(olasDepositBalance);
       },
       0,
@@ -73,7 +82,12 @@ export const MainOlasBalance = () => {
     ]);
 
     return balanceFormat(totalOlasBalance, 2);
-  }, [masterWalletBalances, serviceStakedBalances, serviceWalletBalances]);
+  }, [
+    masterWalletBalances,
+    selectedAgentConfig.requiresAgentSafesOn,
+    serviceStakedBalances,
+    serviceWalletBalances,
+  ]);
 
   return (
     <CardSection

frontend/components/ManageStakingPage/StakingContractSection/CantMigrateAlert.tsx

@@ -1,5 +1,5 @@
 import { Flex, Typography } from 'antd';
-import { isEmpty, isNil } from 'lodash';
+import { isEmpty, isNil, sum } from 'lodash';
 import { useMemo } from 'react';
 
 import { CustomAlert } from '@/components/Alert';
@@ -14,10 +14,7 @@ import {
 } from '@/hooks/useBalanceContext';
 import { useNeedsFunds } from '@/hooks/useNeedsFunds';
 import { useServices } from '@/hooks/useServices';
-import {
-  useActiveStakingContractDetails,
-  useStakingContractContext,
-} from '@/hooks/useStakingContractDetails';
+import { useStakingContractContext } from '@/hooks/useStakingContractDetails';
 import { balanceFormat } from '@/utils/numberFormatters';
 
 import { CantMigrateReason } from './useMigrate';
@@ -32,9 +29,10 @@ const AlertInsufficientMigrationFunds = ({
   const { selectedAgentConfig } = useServices();
   const { isAllStakingContractDetailsRecordLoaded } =
     useStakingContractContext();
-  const { isServiceStaked } = useActiveStakingContractDetails();
-  const { isLoaded: isBalanceLoaded, totalStakedOlasBalance } =
-    useBalanceContext();
+  const {
+    isLoaded: isBalanceLoaded,
+    totalStakedOlasBalance: totalStakedOlasBalanceOnHomeChain,
+  } = useBalanceContext();
   const { masterSafeBalances } = useMasterBalances();
   const { serviceFundRequirements, isInitialFunded } = useNeedsFunds(
     stakingProgramIdToMigrateTo,
@@ -45,7 +43,7 @@ const AlertInsufficientMigrationFunds = ({
       TokenSymbol.OLAS
     ];
 
-  const safeBalance = useMemo(() => {
+  const masterSafeBalanceOnHomeChain = useMemo(() => {
     if (!isBalanceLoaded) return;
     if (isNil(masterSafeBalances) || isEmpty(masterSafeBalances)) return;
     return masterSafeBalances.reduce(
@@ -61,24 +59,26 @@ const AlertInsufficientMigrationFunds = ({
 
   if (!isAllStakingContractDetailsRecordLoaded) return null;
   if (isNil(requiredStakedOlas)) return null;
-  if (isNil(safeBalance?.[TokenSymbol.OLAS])) return null;
+  if (isNil(masterSafeBalanceOnHomeChain?.[TokenSymbol.OLAS])) return null;
 
-  if (isNil(totalStakedOlasBalance)) return null;
+  if (isNil(totalStakedOlasBalanceOnHomeChain)) return null;
 
-  const requiredOlasDeposit = isServiceStaked
-    ? requiredStakedOlas -
-      (totalStakedOlasBalance + safeBalance[TokenSymbol.OLAS]) // when staked
-    : requiredStakedOlas - safeBalance[TokenSymbol.OLAS]; // when not staked
+  const requiredOlasDeposit =
+    requiredStakedOlas -
+    sum([
+      masterSafeBalanceOnHomeChain[TokenSymbol.OLAS],
+      totalStakedOlasBalanceOnHomeChain,
+    ]);
 
   const homeChainId = selectedAgentConfig.evmHomeChainId;
   const nativeTokenSymbol = getNativeTokenSymbol(homeChainId);
   const requiredNativeTokenDeposit = isInitialFunded
     ? selectedAgentConfig.operatingThresholds[WalletOwnerType.Master][
         WalletType.Safe
       ][CHAIN_CONFIG[selectedAgentConfig.evmHomeChainId].nativeToken.symbol] -
-      (safeBalance[nativeTokenSymbol] || 0) // is already funded allow minimal maintenance
+      (masterSafeBalanceOnHomeChain[nativeTokenSymbol] || 0) // is already funded allow minimal maintenance
     : (serviceFundRequirements[homeChainId]?.[nativeTokenSymbol] || 0) -
-      (safeBalance[nativeTokenSymbol] || 0); // otherwise require full initial funding requirements
+      (masterSafeBalanceOnHomeChain[nativeTokenSymbol] || 0); // otherwise require full initial funding requirements
 
   return (
     <CustomAlert

frontend/components/ManageStakingPage/StakingContractSection/MigrateButton.tsx

@@ -3,16 +3,12 @@ import { isNil } from 'lodash';
 import { useMemo } from 'react';
 
 import { MiddlewareDeploymentStatus, ServiceTemplate } from '@/client';
-import {
-  getServiceTemplate,
-  SERVICE_TEMPLATES,
-} from '@/constants/serviceTemplates';
+import { SERVICE_TEMPLATES } from '@/constants/serviceTemplates';
 import { Pages } from '@/enums/Pages';
 import { StakingProgramId } from '@/enums/StakingProgram';
 import { useBalanceContext } from '@/hooks/useBalanceContext';
 import { useModals } from '@/hooks/useModals';
 import { usePageState } from '@/hooks/usePageState';
-import { useService } from '@/hooks/useService';
 import { useServices } from '@/hooks/useServices';
 import {
   useActiveStakingContractDetails,
@@ -43,14 +39,12 @@ export const MigrateButton = ({
     isServicesLoaded && selectedService
       ? selectedService.service_config_id
       : '';
-  const { service } = useService(serviceConfigId);
   const serviceTemplate = useMemo<ServiceTemplate | undefined>(
     () =>
-      (service && getServiceTemplate(service.hash)) ??
       SERVICE_TEMPLATES.find(
         (template) => template.agentType === selectedAgentType,
       ),
-    [selectedAgentType, service],
+    [selectedAgentType],
   );
 
   const { setIsPaused: setIsBalancePollingPaused } = useBalanceContext();

frontend/components/YourWalletPage/index.tsx

@@ -72,14 +72,18 @@ const Address = () => {
 };
 
 const OlasBalance = () => {
+  const { selectedAgentConfig } = useServices();
   const { totalStakedOlasBalance } = useBalanceContext();
   const { masterWalletBalances } = useMasterBalances();
-  const { middlewareChain, evmHomeChainId } = useYourWallet();
+  const { middlewareChain } = useYourWallet();
 
   const masterSafeOlasBalance = masterWalletBalances
     ?.filter(
-      ({ symbol, evmChainId }) =>
-        symbol === TokenSymbol.OLAS && evmChainId === evmHomeChainId,
+      (walletBalance) =>
+        walletBalance.symbol === TokenSymbol.OLAS &&
+        selectedAgentConfig.requiresMasterSafesOn.includes(
+          walletBalance.evmChainId,
+        ),
     )
     .reduce((acc, balance) => acc + balance.balance, 0);