olas-operate-middleware 0.10.20__py3-none-any.whl → 0.11.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -19,18 +19,15 @@
19
19
  # ------------------------------------------------------------------------------
20
20
  """Service manager."""
21
21
 
22
- import asyncio
23
22
  import json
24
23
  import logging
25
24
  import os
26
25
  import traceback
27
26
  import typing as t
28
27
  from collections import Counter, defaultdict
29
- from concurrent.futures import ThreadPoolExecutor
30
28
  from contextlib import suppress
31
29
  from http import HTTPStatus
32
30
  from pathlib import Path
33
- from time import time
34
31
 
35
32
  import requests
36
33
  from aea.helpers.base import IPFSHash
@@ -38,14 +35,16 @@ from aea_ledger_ethereum import LedgerApi
38
35
  from autonomy.chain.base import registry_contracts
39
36
  from autonomy.chain.config import CHAIN_PROFILES, ChainType
40
37
  from autonomy.chain.metadata import IPFS_URI_PREFIX
41
- from web3 import Web3
42
38
 
43
39
  from operate.constants import (
44
40
  AGENT_LOG_DIR,
45
41
  AGENT_LOG_ENV_VAR,
46
42
  AGENT_PERSISTENT_STORAGE_DIR,
47
43
  AGENT_PERSISTENT_STORAGE_ENV_VAR,
44
+ DEFAULT_TOPUP_THRESHOLD,
48
45
  IPFS_ADDRESS,
46
+ MIN_AGENT_BOND,
47
+ MIN_SECURITY_DEPOSIT,
49
48
  ZERO_ADDRESS,
50
49
  )
51
50
  from operate.data import DATA_DIR
@@ -53,21 +52,20 @@ from operate.data.contracts.mech_activity.contract import MechActivityContract
53
52
  from operate.data.contracts.requester_activity_checker.contract import (
54
53
  RequesterActivityCheckerContract,
55
54
  )
56
- from operate.data.contracts.staking_token.contract import StakingTokenContract
57
55
  from operate.keys import KeysManager
58
- from operate.ledger import get_currency_denom, get_default_rpc
56
+ from operate.ledger import get_default_rpc
59
57
  from operate.ledger.profiles import (
60
58
  CONTRACTS,
61
- DEFAULT_MASTER_EOA_FUNDS,
59
+ DEFAULT_EOA_THRESHOLD,
60
+ DEFAULT_EOA_TOPUPS,
62
61
  DEFAULT_PRIORITY_MECH,
63
62
  OLAS,
64
- STAKING,
65
- USDC,
66
63
  WRAPPED_NATIVE_ASSET,
67
64
  get_staking_contract,
68
65
  )
69
66
  from operate.operate_types import (
70
67
  Chain,
68
+ ChainAmounts,
71
69
  FundingValues,
72
70
  LedgerConfig,
73
71
  MechMarketplaceConfig,
@@ -75,7 +73,13 @@ from operate.operate_types import (
75
73
  ServiceEnvProvisionType,
76
74
  ServiceTemplate,
77
75
  )
78
- from operate.services.protocol import EthSafeTxBuilder, OnChainManager, StakingState
76
+ from operate.services.funding_manager import FundingManager
77
+ from operate.services.protocol import (
78
+ EthSafeTxBuilder,
79
+ OnChainManager,
80
+ StakingManager,
81
+ StakingState,
82
+ )
79
83
  from operate.services.service import (
80
84
  ChainConfig,
81
85
  Deployment,
@@ -87,14 +91,16 @@ from operate.services.service import (
87
91
  Service,
88
92
  )
89
93
  from operate.services.utils.mech import deploy_mech
90
- from operate.utils.gnosis import drain_eoa, get_asset_balance, get_assets_balances
91
- from operate.utils.gnosis import transfer as transfer_from_safe
92
- from operate.utils.gnosis import transfer_erc20_from_safe
93
- from operate.wallet.master import MasterWalletManager
94
+ from operate.utils.gnosis import (
95
+ get_asset_balance,
96
+ get_assets_balances,
97
+ transfer_erc20_from_safe,
98
+ )
99
+ from operate.wallet.master import InsufficientFundsException, MasterWalletManager
94
100
 
95
101
 
96
102
  # pylint: disable=redefined-builtin
97
- DEFAULT_TOPUP_THRESHOLD = 0.5
103
+
98
104
  # At the moment, we only support running one agent per service locally on a machine.
99
105
  # If multiple agents are provided in the service.yaml file, only the 0th index config will be used.
100
106
  NUM_LOCAL_AGENT_INSTANCES = 1
@@ -107,6 +113,7 @@ class ServiceManager:
107
113
  self,
108
114
  path: Path,
109
115
  wallet_manager: MasterWalletManager,
116
+ funding_manager: FundingManager,
110
117
  logger: logging.Logger,
111
118
  skip_dependency_check: t.Optional[bool] = False,
112
119
  ) -> None:
@@ -121,6 +128,7 @@ class ServiceManager:
121
128
  self.path = path
122
129
  self.keys_manager = KeysManager()
123
130
  self.wallet_manager = wallet_manager
131
+ self.funding_manager = funding_manager
124
132
  self.logger = logger
125
133
  self.skip_depencency_check = skip_dependency_check
126
134
 
@@ -603,7 +611,7 @@ class ServiceManager:
603
611
  ]
604
612
 
605
613
  except Exception as e: # pylint: disable=broad-except
606
- self.logger.error(f"{e}: {traceback.format_exc()}")
614
+ self.logger.debug(f"{e}: {traceback.format_exc()}")
607
615
  self.logger.warning(
608
616
  "Cannot determine type of activity checker contract. Using default parameters. "
609
617
  "NOTE: This will be an exception in the future!"
@@ -955,7 +963,6 @@ class ServiceManager:
955
963
  self._get_on_chain_state(service=service, chain=chain)
956
964
  == OnChainState.PRE_REGISTRATION
957
965
  ):
958
- # TODO Verify that this is incorrect: cost_of_bond = staking_params["min_staking_deposit"]
959
966
  cost_of_bond = user_params.cost_of_bond
960
967
  if user_params.use_staking:
961
968
  token_utility = target_staking_params["service_registry_token_utility"]
@@ -992,7 +999,7 @@ class ServiceManager:
992
999
  self.logger.info(
993
1000
  f"Approved {token_utility_allowance} OLAS from {safe} to {token_utility}"
994
1001
  )
995
- cost_of_bond = 1
1002
+ cost_of_bond = MIN_AGENT_BOND
996
1003
 
997
1004
  self.logger.info("Activating service")
998
1005
 
@@ -1002,7 +1009,9 @@ class ServiceManager:
1002
1009
  address=safe,
1003
1010
  )
1004
1011
 
1005
- if native_balance < cost_of_bond:
1012
+ if (
1013
+ native_balance < cost_of_bond
1014
+ ): # TODO check that this is the security deposit
1006
1015
  message = f"Cannot activate service: address {safe} {native_balance=} < {cost_of_bond=}."
1007
1016
  self.logger.error(message)
1008
1017
  raise ValueError(message)
@@ -1054,7 +1063,7 @@ class ServiceManager:
1054
1063
  self.logger.info(
1055
1064
  f"Approved {token_utility_allowance} OLAS from {safe} to {token_utility}"
1056
1065
  )
1057
- cost_of_bond = 1 * len(service.agent_addresses)
1066
+ cost_of_bond = MIN_AGENT_BOND
1058
1067
 
1059
1068
  self.logger.info(
1060
1069
  f"Registering agent instances: {chain_data.token} -> {service.agent_addresses}"
@@ -1066,7 +1075,7 @@ class ServiceManager:
1066
1075
  address=safe,
1067
1076
  )
1068
1077
 
1069
- if native_balance < cost_of_bond:
1078
+ if native_balance < cost_of_bond * len(service.agent_addresses):
1070
1079
  message = f"Cannot register agent instances: address {safe} {native_balance=} < {cost_of_bond=}."
1071
1080
  self.logger.error(message)
1072
1081
  raise ValueError(message)
@@ -1081,27 +1090,21 @@ class ServiceManager:
1081
1090
  ).settle()
1082
1091
 
1083
1092
  # Deploy service
1093
+ is_initial_funding = False
1084
1094
  if (
1085
1095
  self._get_on_chain_state(service=service, chain=chain)
1086
1096
  == OnChainState.FINISHED_REGISTRATION
1087
1097
  ):
1088
1098
  self.logger.info("Deploying service")
1089
1099
 
1090
- reuse_multisig = True
1091
1100
  info = sftxb.info(token_id=chain_data.token)
1092
1101
  service_safe_address = info["multisig"]
1093
1102
  if service_safe_address == ZERO_ADDRESS:
1094
1103
  reuse_multisig = False
1095
-
1096
- self.logger.info(f"{reuse_multisig=}")
1097
-
1098
- is_recovery_module_enabled = (
1099
- True # Ensure is true for non-deployed multisigs
1100
- )
1101
- if (
1102
- service_safe_address is not None
1103
- and service_safe_address != ZERO_ADDRESS
1104
- ):
1104
+ is_initial_funding = True
1105
+ is_recovery_module_enabled = True
1106
+ else:
1107
+ reuse_multisig = True
1105
1108
  is_recovery_module_enabled = (
1106
1109
  registry_contracts.gnosis_safe.is_module_enabled(
1107
1110
  ledger_api=sftxb.ledger_api,
@@ -1110,6 +1113,7 @@ class ServiceManager:
1110
1113
  ).get("enabled")
1111
1114
  )
1112
1115
 
1116
+ self.logger.info(f"{reuse_multisig=}")
1113
1117
  self.logger.info(f"{is_recovery_module_enabled=}")
1114
1118
 
1115
1119
  messages = sftxb.get_deploy_data_from_safe(
@@ -1128,6 +1132,9 @@ class ServiceManager:
1128
1132
  chain_data.instances = info["instances"]
1129
1133
  chain_data.multisig = info["multisig"]
1130
1134
 
1135
+ if is_initial_funding:
1136
+ self.funding_manager.fund_service_initial(service)
1137
+
1131
1138
  # TODO: yet another agent specific logic for mech, which should be abstracted
1132
1139
  if all(
1133
1140
  var in service.env_variables
@@ -1239,7 +1246,6 @@ class ServiceManager:
1239
1246
  self,
1240
1247
  service_config_id: str,
1241
1248
  chain: str,
1242
- withdrawal_address: t.Optional[str] = None,
1243
1249
  ) -> None:
1244
1250
  """Terminate service on-chain"""
1245
1251
 
@@ -1249,10 +1255,7 @@ class ServiceManager:
1249
1255
  ledger_config = chain_config.ledger_config
1250
1256
  chain_data = chain_config.chain_data
1251
1257
  wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
1252
- safe = wallet.safes[Chain(chain)] # type: ignore
1253
-
1254
- if withdrawal_address:
1255
- withdrawal_address = Web3.to_checksum_address(withdrawal_address)
1258
+ master_safe = wallet.safes[Chain(chain)] # type: ignore
1256
1259
 
1257
1260
  # TODO fixme
1258
1261
  os.environ["CUSTOM_CHAIN_RPC"] = ledger_config.rpc
@@ -1277,23 +1280,23 @@ class ServiceManager:
1277
1280
  )
1278
1281
 
1279
1282
  # Cannot unstake, terminate flow.
1280
- if is_staked and not can_unstake and withdrawal_address is None:
1283
+ if is_staked and not can_unstake:
1281
1284
  self.logger.info("Service cannot be terminated on-chain: cannot unstake.")
1282
1285
  return
1283
-
1284
1286
  # Unstake the service if applies
1285
- if is_staked and (can_unstake or withdrawal_address is not None):
1287
+ if is_staked and can_unstake:
1286
1288
  self.unstake_service_on_chain_from_safe(
1287
1289
  service_config_id=service_config_id,
1288
1290
  chain=chain,
1289
1291
  staking_program_id=current_staking_program,
1290
1292
  )
1293
+ # At least claim the rewards if we cannot unstake yet
1291
1294
  elif is_staked:
1292
- # at least claim the rewards if we cannot unstake yet
1293
1295
  self.claim_on_chain_from_safe(
1294
1296
  service_config_id=service_config_id,
1295
1297
  chain=chain,
1296
1298
  )
1299
+ return
1297
1300
 
1298
1301
  if self._get_on_chain_state(service=service, chain=chain) in (
1299
1302
  OnChainState.ACTIVE_REGISTRATION,
@@ -1323,38 +1326,58 @@ class ServiceManager:
1323
1326
  counter_current_safe_owners = Counter(s.lower() for s in current_safe_owners)
1324
1327
  counter_instances = Counter(s.lower() for s in service.agent_addresses)
1325
1328
 
1326
- if withdrawal_address is not None:
1327
- # we don't drain signer yet, because the owner swapping tx may need to happen
1328
- self.drain_service_safe(
1329
- service_config_id=service_config_id,
1330
- withdrawal_address=withdrawal_address,
1331
- chain=Chain(chain),
1332
- )
1333
-
1334
1329
  if counter_current_safe_owners == counter_instances:
1335
- if withdrawal_address is None:
1336
- self.logger.info("Service funded for safe swap")
1337
- self.fund_service(
1338
- service_config_id=service_config_id,
1339
- funding_values={
1340
- ZERO_ADDRESS: {
1341
- "agent": {
1342
- "topup": chain_data.user_params.fund_requirements[
1343
- ZERO_ADDRESS
1344
- ].agent,
1345
- "threshold": chain_data.user_params.fund_requirements[
1346
- ZERO_ADDRESS
1347
- ].agent,
1348
- },
1349
- "safe": {"topup": 0, "threshold": 0},
1330
+ requirements = ChainAmounts(
1331
+ {
1332
+ chain: {
1333
+ current_safe_owners[0]: {
1334
+ ZERO_ADDRESS: chain_data.user_params.fund_requirements[
1335
+ ZERO_ADDRESS
1336
+ ].agent
1350
1337
  }
1351
- },
1338
+ }
1339
+ }
1340
+ )
1341
+ balances = ChainAmounts(
1342
+ {
1343
+ chain: {
1344
+ current_safe_owners[0]: {
1345
+ ZERO_ADDRESS: get_asset_balance(
1346
+ ledger_api=sftxb.ledger_api,
1347
+ asset_address=ZERO_ADDRESS,
1348
+ address=service.agent_addresses[0],
1349
+ )
1350
+ }
1351
+ }
1352
+ }
1353
+ )
1354
+ if balances < requirements * DEFAULT_EOA_THRESHOLD:
1355
+ self.logger.info("[SERVICE MANAGER] Funding agent EOA for Safe swap.")
1356
+ shortfalls = ChainAmounts.shortfalls(
1357
+ requirements=requirements, balances=balances
1352
1358
  )
1359
+ try:
1360
+ self.funding_manager.fund_chain_amounts(shortfalls)
1361
+ except InsufficientFundsException as e:
1362
+ recovery_module_address = CONTRACTS[Chain(chain)]["recovery_module"]
1363
+ is_recovery_module_enabled = (
1364
+ registry_contracts.gnosis_safe.is_module_enabled(
1365
+ ledger_api=sftxb.ledger_api,
1366
+ contract_address=chain_data.multisig,
1367
+ module_address=recovery_module_address,
1368
+ ).get("enabled")
1369
+ )
1370
+ if is_recovery_module_enabled:
1371
+ self.logger.info(
1372
+ "[SERVICE MANAGER] Could not fund Agent EOA for service swap, but recovery module is enabled."
1373
+ )
1374
+ return
1375
+ raise e
1353
1376
 
1354
1377
  self._enable_recovery_module(
1355
1378
  service_config_id=service_config_id, chain=chain
1356
1379
  )
1357
- self.logger.info("Swapping Safe owners")
1380
+ self.logger.info("[SERVICE MANAGER] Swapping Safe owners")
1358
1381
  owner_crypto = self.keys_manager.get_crypto_instance(
1359
1382
  address=current_safe_owners[0]
1360
1383
  )
@@ -1363,25 +1386,10 @@ class ServiceManager:
1363
1386
  multisig=chain_data.multisig, # TODO this can be read from the registry
1364
1387
  owner_cryptos=[owner_crypto], # TODO allow multiple owners
1365
1388
  new_owner_address=(
1366
- safe if safe else wallet.crypto.address
1389
+ master_safe if master_safe else wallet.crypto.address
1367
1390
  ), # TODO it should always be safe address
1368
1391
  )
1369
1392
 
1370
- if withdrawal_address is not None:
1371
- ethereum_crypto = KeysManager().get_crypto_instance(
1372
- service.agent_addresses[0]
1373
- )
1374
- # drain all native tokens from service signer key
1375
- drain_eoa(
1376
- ledger_api=self.wallet_manager.load(
1377
- ledger_config.chain.ledger_type
1378
- ).ledger_api(chain=ledger_config.chain, rpc=ledger_config.rpc),
1379
- crypto=ethereum_crypto,
1380
- withdrawal_address=withdrawal_address,
1381
- chain_id=ledger_config.chain.id,
1382
- )
1383
- self.logger.info(f"{service.name} signer drained")
1384
-
1385
1393
  def _execute_recovery_module_flow_from_safe( # pylint: disable=too-many-locals
1386
1394
  self,
1387
1395
  service_config_id: str,
@@ -1524,82 +1532,14 @@ class ServiceManager:
1524
1532
  f"Cannot enable recovery module. Safe {service_safe_address} has inconsistent owners."
1525
1533
  )
1526
1534
 
1527
- def _get_current_staking_program(
1535
+ def _get_current_staking_program( # pylint: disable=no-self-use
1528
1536
  self, service: Service, chain: str
1529
1537
  ) -> t.Optional[str]:
1530
- chain_config = service.chain_configs[chain]
1531
- ledger_config = chain_config.ledger_config
1532
- sftxb = self.get_eth_safe_tx_builder(ledger_config=ledger_config)
1533
- service_id = chain_config.chain_data.token
1534
- ledger_api = sftxb.ledger_api
1535
-
1536
- if service_id == NON_EXISTENT_TOKEN:
1537
- return None
1538
-
1539
- service_registry = registry_contracts.service_registry.get_instance(
1540
- ledger_api=ledger_api,
1541
- contract_address=CONTRACTS[ledger_config.chain]["service_registry"],
1542
- )
1543
-
1544
- service_owner = service_registry.functions.ownerOf(service_id).call()
1545
-
1546
- # TODO Implement in Staking Manager. Implemented here for performance issues.
1547
- staking_ctr = t.cast(
1548
- StakingTokenContract,
1549
- StakingTokenContract.from_dir(
1550
- directory=str(DATA_DIR / "contracts" / "staking_token")
1551
- ),
1538
+ staking_manager = StakingManager(Chain(chain))
1539
+ return staking_manager.get_current_staking_program(
1540
+ service_id=service.chain_configs[chain].chain_data.token
1552
1541
  )
1553
1542
 
1554
- try:
1555
- state = StakingState(
1556
- staking_ctr.get_instance(
1557
- ledger_api=ledger_api,
1558
- contract_address=service_owner,
1559
- )
1560
- .functions.getStakingState(service_id)
1561
- .call()
1562
- )
1563
- except Exception: # pylint: disable=broad-except
1564
- # Service owner is not a staking contract
1565
-
1566
- # TODO The exception caught here should be ContractLogicError.
1567
- # This exception is typically raised when the contract reverts with
1568
- # a reason string. However, in some cases, the error message
1569
- # does not contain a reason string, which means web3.py raises
1570
- # a generic ValueError instead. It should be properly analyzed
1571
- # what exceptions might be raised by web3.py in this case. To
1572
- # avoid any issues we are simply catching all exceptions.
1573
- return None
1574
-
1575
- if state == StakingState.UNSTAKED:
1576
- return None
1577
-
1578
- for staking_program_id, val in STAKING[ledger_config.chain].items():
1579
- if val == service_owner:
1580
- return staking_program_id
1581
-
1582
- # Fallback, if not possible to determine staking_program_id it means it's an "inner" staking contract
1583
- # (e.g., in the case of DualStakingToken). Loop trough all the known contracts.
1584
- for staking_program_id, staking_program_address in STAKING[
1585
- ledger_config.chain
1586
- ].items():
1587
- state = StakingState(
1588
- staking_ctr.get_instance(
1589
- ledger_api=ledger_api,
1590
- contract_address=staking_program_address,
1591
- )
1592
- .functions.getStakingState(service_id)
1593
- .call()
1594
- )
1595
-
1596
- if state in (StakingState.STAKED, StakingState.EVICTED):
1597
- return staking_program_id
1598
-
1599
- # it's staked, but we don't know which staking program
1600
- # so the staking_program_id should be an arbitrary staking contract
1601
- return service_owner
1602
-
1603
1543
  def unbond_service_on_chain(
1604
1544
  self, service_config_id: str, chain: t.Optional[str] = None
1605
1545
  ) -> None:
@@ -1906,6 +1846,16 @@ class ServiceManager:
1906
1846
  )
1907
1847
  ).settle()
1908
1848
 
1849
+ def claim_all_on_chain_from_safe(self) -> None:
1850
+ """Claim rewards from all services and chains"""
1851
+ self.logger.info("claim_all_on_chain_from_safe")
1852
+ services, _ = self.get_all_services()
1853
+ for service in services:
1854
+ self.claim_on_chain_from_safe(
1855
+ service_config_id=service.service_config_id,
1856
+ chain=service.home_chain,
1857
+ )
1858
+
1909
1859
  def claim_on_chain_from_safe(
1910
1860
  self,
1911
1861
  service_config_id: str,
@@ -1918,6 +1868,14 @@ class ServiceManager:
1918
1868
  ledger_config = chain_config.ledger_config
1919
1869
  wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
1920
1870
  ledger_api = wallet.ledger_api(chain=ledger_config.chain, rpc=ledger_config.rpc)
1871
+
1872
+ if (
1873
+ chain_config.chain_data.token == NON_EXISTENT_TOKEN
1874
+ or chain_config.chain_data.multisig == ZERO_ADDRESS
1875
+ ):
1876
+ self.logger.info("Service is not minted or Safe not deployed.")
1877
+ return 0
1878
+
1921
1879
  self.logger.info(
1922
1880
  f"OLAS Balance on service Safe {chain_config.chain_data.multisig}: "
1923
1881
  f"{get_asset_balance(ledger_api, OLAS[Chain(chain)], chain_config.chain_data.multisig)}"
@@ -1930,10 +1888,11 @@ class ServiceManager:
1930
1888
  staking_program_id=current_staking_program,
1931
1889
  )
1932
1890
  if staking_contract is None:
1933
- raise RuntimeError(
1891
+ self.logger.warning(
1934
1892
  "No staking contract found for the "
1935
1893
  f"{current_staking_program=}. Not claiming the rewards."
1936
1894
  )
1895
+ return 0
1937
1896
 
1938
1897
  sftxb = self.get_eth_safe_tx_builder(ledger_config=ledger_config)
1939
1898
  if not sftxb.staking_rewards_claimable(
@@ -1978,22 +1937,13 @@ class ServiceManager:
1978
1937
  def fund_service( # pylint: disable=too-many-arguments,too-many-locals
1979
1938
  self,
1980
1939
  service_config_id: str,
1981
- funding_values: t.Optional[FundingValues] = None,
1982
- from_safe: bool = True,
1983
- task_id: t.Optional[str] = None,
1940
+ amounts: ChainAmounts,
1984
1941
  ) -> None:
1985
1942
  """Fund service if required."""
1986
1943
  service = self.load(service_config_id=service_config_id)
1944
+ self.funding_manager.fund_service(service=service, amounts=amounts)
1987
1945
 
1988
- for chain in service.chain_configs.keys():
1989
- self.logger.info(f"[FUNDING_JOB] [{task_id=}] Funding {chain=}")
1990
- self.fund_service_single_chain(
1991
- service_config_id=service_config_id,
1992
- funding_values=funding_values,
1993
- from_safe=from_safe,
1994
- chain=chain,
1995
- )
1996
-
1946
+ # TODO deprecate
1997
1947
  def fund_service_single_chain( # pylint: disable=too-many-arguments,too-many-locals,too-many-statements
1998
1948
  self,
1999
1949
  service_config_id: str,
@@ -2026,7 +1976,10 @@ class ServiceManager:
2026
1976
  else:
2027
1977
  on_chain_operations_buffer = (
2028
1978
  chain_data.user_params.cost_of_bond
2029
- * (1 + len(service.agent_addresses))
1979
+ * (
1980
+ MIN_SECURITY_DEPOSIT
1981
+ + MIN_AGENT_BOND * len(service.agent_addresses)
1982
+ )
2030
1983
  )
2031
1984
 
2032
1985
  asset_funding_values = (
@@ -2071,10 +2024,13 @@ class ServiceManager:
2071
2024
  to_transfer = max(
2072
2025
  min(available_balance, target_balance - agent_balance), 0
2073
2026
  )
2027
+ if to_transfer <= 0:
2028
+ continue
2029
+
2074
2030
  self.logger.info(
2075
2031
  f"[FUNDING_JOB] Transferring {to_transfer} units (asset {asset_address}) to agent {agent_address}"
2076
2032
  )
2077
- wallet.transfer_asset(
2033
+ wallet.transfer(
2078
2034
  asset=asset_address,
2079
2035
  to=agent_address,
2080
2036
  amount=int(to_transfer),
@@ -2143,7 +2099,7 @@ class ServiceManager:
2143
2099
  # when not enough funds are present, and the FE doesn't let the user to start the agent.
2144
2100
  # Ideally this error should be allowed, and then the FE should ask the user for more funds.
2145
2101
  with suppress(RuntimeError):
2146
- wallet.transfer_asset(
2102
+ wallet.transfer(
2147
2103
  asset=asset_address,
2148
2104
  to=t.cast(str, chain_data.multisig),
2149
2105
  amount=int(to_transfer),
@@ -2151,6 +2107,7 @@ class ServiceManager:
2151
2107
  rpc=rpc or ledger_config.rpc,
2152
2108
  )
2153
2109
 
2110
+ # TODO Deprecate
2154
2111
  # TODO This method is possibly not used anymore
2155
2112
  def fund_service_erc20( # pylint: disable=too-many-arguments,too-many-locals
2156
2113
  self,
@@ -2188,9 +2145,12 @@ class ServiceManager:
2188
2145
  agent_topup
2189
2146
  or chain_data.user_params.fund_requirements[ZERO_ADDRESS].agent
2190
2147
  )
2148
+ if to_transfer <= 0:
2149
+ continue
2150
+
2191
2151
  self.logger.info(f"Transferring {to_transfer} units to {agent_address}")
2192
- wallet.transfer_erc20(
2193
- token=token,
2152
+ wallet.transfer(
2153
+ asset=token,
2194
2154
  to=agent_address,
2195
2155
  amount=int(to_transfer),
2196
2156
  chain=ledger_config.chain,
@@ -2215,150 +2175,36 @@ class ServiceManager:
2215
2175
  safe_topup
2216
2176
  or chain_data.user_params.fund_requirements[ZERO_ADDRESS].safe
2217
2177
  )
2178
+ if to_transfer <= 0:
2179
+ return
2180
+
2218
2181
  self.logger.info(
2219
2182
  f"Transferring {to_transfer} units to {chain_data.multisig}"
2220
2183
  )
2221
- wallet.transfer_erc20(
2222
- token=token,
2184
+ wallet.transfer(
2185
+ asset=token,
2223
2186
  to=t.cast(str, chain_data.multisig),
2224
2187
  amount=int(to_transfer),
2225
2188
  chain=ledger_config.chain,
2226
2189
  rpc=rpc or ledger_config.rpc,
2227
2190
  )
2228
2191
 
2229
- def drain_service_safe( # pylint: disable=too-many-locals
2230
- self,
2231
- service_config_id: str,
2232
- withdrawal_address: str,
2233
- chain: Chain,
2234
- ) -> None:
2235
- """Drain the funds out of the service safe."""
2236
- self.logger.info(
2237
- f"Draining the safe of service: {service_config_id} on chain {chain.value}"
2238
- )
2239
- service = self.load(service_config_id=service_config_id)
2240
- chain_config = service.chain_configs[chain.value]
2241
- ledger_config = chain_config.ledger_config
2242
- chain_data = chain_config.chain_data
2243
- wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
2244
- ledger_api = wallet.ledger_api(chain=ledger_config.chain, rpc=ledger_config.rpc)
2245
- ethereum_crypto = KeysManager().get_crypto_instance(service.agent_addresses[0])
2246
- withdrawal_address = Web3.to_checksum_address(withdrawal_address)
2247
-
2248
- # drain ERC20 tokens from service safe
2249
- for token_name, token_address in (
2250
- ("OLAS", OLAS[chain]),
2251
- (
2252
- f"W{get_currency_denom(chain)}",
2253
- WRAPPED_NATIVE_ASSET[chain],
2254
- ),
2255
- ("USDC", USDC[chain]),
2256
- ):
2257
- token_instance = registry_contracts.erc20.get_instance(
2258
- ledger_api=ledger_api,
2259
- contract_address=token_address,
2260
- )
2261
- balance = token_instance.functions.balanceOf(chain_data.multisig).call()
2262
- if balance == 0:
2263
- self.logger.info(
2264
- f"No {token_name} to drain from service safe: {chain_data.multisig}"
2265
- )
2266
- continue
2267
-
2268
- self.logger.info(
2269
- f"Draining {balance} {token_name} out of service safe: {chain_data.multisig}"
2270
- )
2271
- transfer_erc20_from_safe(
2272
- ledger_api=ledger_api,
2273
- crypto=ethereum_crypto,
2274
- safe=chain_data.multisig,
2275
- token=token_address,
2276
- to=withdrawal_address,
2277
- amount=balance,
2278
- )
2279
-
2280
- # drain native asset from service safe
2281
- balance = ledger_api.get_balance(chain_data.multisig)
2282
- if balance == 0:
2283
- self.logger.info(
2284
- f"No {get_currency_denom(chain)} to drain from service safe: {chain_data.multisig}"
2285
- )
2286
- else:
2287
- self.logger.info(
2288
- f"Draining {balance} {get_currency_denom(chain)} out of service safe: {chain_data.multisig}"
2289
- )
2290
- transfer_from_safe(
2291
- ledger_api=ledger_api,
2292
- crypto=ethereum_crypto,
2293
- safe=chain_data.multisig,
2294
- to=withdrawal_address,
2295
- amount=balance,
2296
- )
2297
-
2298
- self.logger.info(f"{service.name} safe drained ({service_config_id=})")
2299
-
2300
- async def funding_job(
2301
- self,
2302
- service_config_id: str,
2303
- loop: t.Optional[asyncio.AbstractEventLoop] = None,
2304
- from_safe: bool = True,
2192
+ def drain(
2193
+ self, service_config_id: str, chain_str: str, withdrawal_address: str
2305
2194
  ) -> None:
2306
- """Start a background funding job."""
2307
- loop = loop or asyncio.get_event_loop()
2195
+ """Drain service safe and agent EOAs."""
2308
2196
  service = self.load(service_config_id=service_config_id)
2309
- chain_config = service.chain_configs[service.home_chain]
2310
- task = asyncio.current_task()
2311
- task_id = id(task) if task else "Unknown task_id"
2312
- with ThreadPoolExecutor() as executor:
2313
- last_claim = 0
2314
- while True:
2315
- try:
2316
- await loop.run_in_executor(
2317
- executor,
2318
- self.fund_service,
2319
- service_config_id, # Service id
2320
- {
2321
- asset_address: {
2322
- "agent": {
2323
- "topup": fund_requirements.agent,
2324
- "threshold": int(
2325
- fund_requirements.agent
2326
- * DEFAULT_TOPUP_THRESHOLD
2327
- ),
2328
- },
2329
- "safe": {
2330
- "topup": fund_requirements.safe,
2331
- "threshold": int(
2332
- fund_requirements.safe * DEFAULT_TOPUP_THRESHOLD
2333
- ),
2334
- },
2335
- }
2336
- for asset_address, fund_requirements in chain_config.chain_data.user_params.fund_requirements.items()
2337
- },
2338
- from_safe,
2339
- task_id,
2340
- )
2341
- except Exception: # pylint: disable=broad-except
2342
- logging.info(
2343
- f"Error occured while funding the service\n{traceback.format_exc()}"
2344
- )
2345
-
2346
- # try claiming rewards every hour
2347
- if last_claim + 3600 < time():
2348
- try:
2349
- await loop.run_in_executor(
2350
- executor,
2351
- self.claim_on_chain_from_safe,
2352
- service_config_id,
2353
- service.home_chain,
2354
- )
2355
- except Exception: # pylint: disable=broad-except
2356
- logging.info(
2357
- f"Error occured while claiming rewards\n{traceback.format_exc()}"
2358
- )
2359
- last_claim = time()
2360
-
2361
- await asyncio.sleep(60)
2197
+ chain = Chain(chain_str)
2198
+ self.funding_manager.drain_service_safe(
2199
+ service=service,
2200
+ withdrawal_address=withdrawal_address,
2201
+ chain=chain,
2202
+ )
2203
+ self.funding_manager.drain_agents_eoas(
2204
+ service=service,
2205
+ withdrawal_address=withdrawal_address,
2206
+ chain=chain,
2207
+ )
2362
2208
 
2363
2209
  def deploy_service_locally(
2364
2210
  self,
@@ -2432,6 +2278,14 @@ class ServiceManager:
2432
2278
  )
2433
2279
  return service
2434
2280
 
2281
+ def funding_requirements( # pylint: disable=too-many-locals,too-many-statements,too-many-nested-blocks
2282
+ self, service_config_id: str
2283
+ ) -> t.Dict:
2284
+ """Get the funding requirements for a service."""
2285
+ service = self.load(service_config_id=service_config_id)
2286
+ return self.funding_manager.funding_requirements(service)
2287
+
2288
+ # TODO deprecate
2435
2289
  def refill_requirements( # pylint: disable=too-many-locals,too-many-statements,too-many-nested-blocks
2436
2290
  self, service_config_id: str
2437
2291
  ) -> t.Dict:
@@ -2635,6 +2489,7 @@ class ServiceManager:
2635
2489
  "allow_start_agent": allow_start_agent,
2636
2490
  }
2637
2491
 
2492
+ # TODO deprecate
2638
2493
  def _compute_bonded_assets( # pylint: disable=too-many-locals
2639
2494
  self, service_config_id: str, chain: str
2640
2495
  ) -> t.Dict:
@@ -2750,6 +2605,7 @@ class ServiceManager:
2750
2605
 
2751
2606
  return dict(bonded_assets)
2752
2607
 
2608
+ # TODO deprecate
2753
2609
  def _compute_protocol_asset_requirements( # pylint: disable=too-many-locals
2754
2610
  self, service_config_id: str, chain: str
2755
2611
  ) -> t.Dict:
@@ -2782,6 +2638,7 @@ class ServiceManager:
2782
2638
  ),
2783
2639
  )
2784
2640
 
2641
+ # TODO address this comment in FundingManager
2785
2642
  # This computation assumes the service will be/has been minted with these
2786
2643
  # parameters. Otherwise, these values should be retrieved on-chain as follows:
2787
2644
  # - agent_bonds: by combining the output of ServiceRegistry .getAgentParams .getService
@@ -2797,6 +2654,7 @@ class ServiceManager:
2797
2654
 
2798
2655
  return dict(service_asset_requirements)
2799
2656
 
2657
+ # TODO deprecate
2800
2658
  @staticmethod
2801
2659
  def _compute_refill_requirement(
2802
2660
  asset_funding_values: t.Dict,
@@ -2877,12 +2735,13 @@ class ServiceManager:
2877
2735
  "recommended_refill": recommended_refill,
2878
2736
  }
2879
2737
 
2738
+ # TODO deprecate
2880
2739
  @staticmethod
2881
2740
  def get_master_eoa_native_funding_values(
2882
2741
  master_safe_exists: bool, chain: Chain, balance: int
2883
2742
  ) -> t.Dict:
2884
2743
  """Get Master EOA native funding values."""
2885
2744
 
2886
- topup = DEFAULT_MASTER_EOA_FUNDS[chain][ZERO_ADDRESS]
2745
+ topup = DEFAULT_EOA_TOPUPS[chain][ZERO_ADDRESS]
2887
2746
  threshold = topup / 2 if master_safe_exists else topup
2888
2747
  return {"topup": topup, "threshold": threshold, "balance": balance}