feat: integrated name service contract

src/uagents/agent.py

@@ -4,6 +4,10 @@
 
 from cosmpy.aerial.wallet import LocalWallet, PrivateKey
 from cosmpy.crypto.address import Address
+from cosmpy.aerial.contract.cosmwasm import create_cosmwasm_execute_msg
+from cosmpy.aerial.client import prepare_and_broadcast_basic_transaction
+from cosmpy.aerial.tx import Transaction
+
 
 from uagents.asgi import ASGIServer
 from uagents.context import (
@@ -19,13 +23,19 @@
 from uagents.protocol import Protocol
 from uagents.resolver import Resolver, AlmanacResolver
 from uagents.storage import KeyValueStore, get_or_create_private_keys
-from uagents.network import get_ledger, get_reg_contract, wait_for_tx_to_complete
+from uagents.network import (
+    get_ledger,
+    get_almanac_contract,
+    get_service_contract,
+    wait_for_tx_to_complete,
+)
 from uagents.mailbox import MailboxClient
 from uagents.config import (
+    CONTRACT_ALMANAC,
+    CONTRACT_SERVICE,
     REGISTRATION_FEE,
     REGISTRATION_DENOM,
     LEDGER_PREFIX,
-    BLOCK_INTERVAL,
     parse_endpoint_config,
     parse_mailbox_config,
     get_logger,
@@ -99,7 +109,8 @@ def __init__(
             ]
 
         self._ledger = get_ledger()
-        self._reg_contract = get_reg_contract()
+        self._almanac_contract = get_almanac_contract()
+        self._service_contract = get_service_contract()
         self._storage = KeyValueStore(self.address[0:16])
         self._interval_handlers: List[Tuple[IntervalCallback, float]] = []
         self._interval_messages: Set[str] = set()
@@ -170,9 +181,10 @@ def sign_digest(self, digest: bytes) -> str:
         return self._identity.sign_digest(digest)
 
     def sign_registration(self) -> str:
-        assert self._reg_contract.address is not None
+        assert self._almanac_contract.address is not None
         return self._identity.sign_registration(
-            str(self._reg_contract.address), self.get_registration_sequence()
+            str(self._almanac_contract.address),
+            self._almanac_contract.get_sequence(self.address),
         )
 
     def update_loop(self, loop):
@@ -193,50 +205,66 @@ async def _register(self, ctx: Context):
 
         signature = self.sign_registration()
 
-        msg = {
-            "register": {
-                "record": {
-                    "service": {
-                        "protocols": list(
-                            map(lambda x: x.digest, self.protocols.values())
-                        ),
-                        "endpoints": self._endpoints,
-                    }
-                },
-                "signature": signature,
-                "sequence": self.get_registration_sequence(),
-                "agent_address": self.address,
-            }
-        }
-
-        self._logger.info("Registering on Almanac contract...")
-        transaction = self._reg_contract.execute(
-            msg,
-            ctx.wallet,
-            funds=f"{REGISTRATION_FEE}{REGISTRATION_DENOM}",
-        )
-        await wait_for_tx_to_complete(transaction.tx_hash)
-        self._logger.info("Registering on Almanac contract...complete")
+        self._logger.info("Registering on contract...")
 
-    def _schedule_registration(self):
-        query_msg = {"query_records": {"agent_address": self.address}}
-        response = self._reg_contract.query(query_msg)
+        if not self._service_contract.is_name_available(
+            ctx.name
+        ) and not self._service_contract.is_owner(ctx.name, str(self.wallet.address())):
+            self._logger.error(
+                f"Please select another name for your agent, {ctx.name} is owned by another address"
+            )
+            return
+
+        transaction = Transaction()
 
-        if not response["record"]:
-            contract_state = self._reg_contract.query({"query_contract_state": {}})
-            expiry = contract_state.get("state").get("expiry_height")
-            return expiry * BLOCK_INTERVAL
+        almanac_msg = self._almanac_contract.get_registration_msg(
+            self.protocols, self._endpoints, signature, self.address
+        )
+        ownership_msg = self._service_contract.get_ownership_msg(
+            ctx.name, str(self.wallet.address())
+        )
+        registration_msg = self._service_contract.get_registration_msg(
+            ctx.name, self.address
+        )
 
-        expiry = response.get("record")[0].get("expiry")
-        height = response.get("height")
+        transaction.add_message(
+            create_cosmwasm_execute_msg(
+                ctx.wallet.address(),
+                CONTRACT_ALMANAC,
+                almanac_msg,
+                funds=f"{REGISTRATION_FEE}{REGISTRATION_DENOM}",
+            )
+        )
 
-        return (expiry - height) * BLOCK_INTERVAL
+        transaction.add_message(
+            create_cosmwasm_execute_msg(
+                ctx.wallet.address(), CONTRACT_SERVICE, ownership_msg
+            )
+        )
+        transaction.add_message(
+            create_cosmwasm_execute_msg(
+                ctx.wallet.address(), CONTRACT_SERVICE, registration_msg
+            )
+        )
 
-    def get_registration_sequence(self) -> int:
-        query_msg = {"query_sequence": {"agent_address": self.address}}
-        sequence = self._reg_contract.query(query_msg)["sequence"]
+        transaction = prepare_and_broadcast_basic_transaction(
+            ctx.ledger, transaction, ctx.wallet
+        )
+        await wait_for_tx_to_complete(transaction.tx_hash)
+        self._logger.info("Registering on contract...complete")
 
-        return sequence
+    def _schedule_registration(self):
+        return self._almanac_contract.get_expiry(self.address)
+
+    def get_agent_address(self, name: str) -> str:
+        query_msg = {"domain_record": {"domain": f"{name}.agent"}}
+        res = self._service_contract.query(query_msg)
+        if res["record"] is not None:
+            registered_address = res["record"]["records"][0]["agent_address"]["records"]
+            if len(registered_address) > 0:
+                return registered_address[0]["address"]
+            return 0
+        return 1
 
     def on_interval(
         self,
@@ -343,9 +371,17 @@ def start_background_tasks(self):
 
         # start the contract registration update loop
         if self._endpoints is not None:
-            self._loop.create_task(
-                _run_interval(self._register, self._ctx, self._schedule_registration())
-            )
+            if (
+                not self._almanac_contract.is_registered(self.address)
+                or self._schedule_registration() < 3600
+            ):
+                self._loop.create_task(
+                    _run_interval(
+                        self._register, self._ctx, self._schedule_registration()
+                    )
+                )
+            else:
+                self._logger.info("Registration up to date!")
         else:
             self._logger.warning(
                 "I have no endpoint and won't be able to receive external messages"

src/uagents/config.py

@@ -16,7 +16,8 @@ class AgentNetwork(Enum):
 AGENT_PREFIX = "agent"
 LEDGER_PREFIX = "fetch"
 USER_PREFIX = "user"
-CONTRACT_ALMANAC = "fetch1tjagw8g8nn4cwuw00cf0m5tl4l6wfw9c0ue507fhx9e3yrsck8zs0l3q4w"
+CONTRACT_ALMANAC = "fetch1h5rhtj5m6dqjmufj5m3t4mq6l7cnd8dvaxclwmrk6tfdm0gy3lmszksf0s"
+CONTRACT_SERVICE = "fetch1yrf4xpglq02fzj50m9wn44qdq89a5vr0ufa42qa506uhwal4n79s99sp87"
 REGISTRATION_FEE = 500000000000000000
 REGISTRATION_DENOM = "atestfet"
 BLOCK_INTERVAL = 5

src/uagents/network.py

@@ -1,6 +1,6 @@
 import asyncio
 from datetime import datetime, timedelta
-from typing import Optional
+from typing import Optional, Dict, List, Union
 
 from cosmpy.aerial.contract import LedgerContract
 from cosmpy.aerial.client import (
@@ -13,7 +13,101 @@
 from cosmpy.aerial.faucet import FaucetApi
 from cosmpy.aerial.tx_helpers import TxResponse
 
-from uagents.config import AgentNetwork, CONTRACT_ALMANAC, AGENT_NETWORK
+from uagents.config import (
+    AgentNetwork,
+    CONTRACT_ALMANAC,
+    CONTRACT_SERVICE,
+    AGENT_NETWORK,
+    BLOCK_INTERVAL,
+)
+
+
+class AlmanacContract(LedgerContract):
+    def is_registered(self, address: str) -> bool:
+        query_msg = {"query_records": {"agent_address": address}}
+        response = self.query(query_msg)
+
+        if not response["record"]:
+            return False
+        return True
+
+    def get_expiry(self, address: str):
+        query_msg = {"query_records": {"agent_address": address}}
+        response = self.query(query_msg)
+
+        if not response["record"]:
+            contract_state = self.query({"query_contract_state": {}})
+            expiry = contract_state.get("state").get("expiry_height")
+            return expiry * BLOCK_INTERVAL
+
+        expiry = response.get("record")[0].get("expiry")
+        height = response.get("height")
+
+        return (expiry - height) * BLOCK_INTERVAL
+
+    def get_registration_msg(
+        self,
+        protocols: Dict,
+        endpoints: Optional[Union[List[str], Dict[str, dict]]],
+        signature: str,
+        address: str,
+    ) -> dict:
+        return {
+            "register": {
+                "record": {
+                    "service": {
+                        "protocols": list(map(lambda x: x.digest, protocols.values())),
+                        "endpoints": endpoints,
+                    }
+                },
+                "signature": signature,
+                "sequence": self.get_sequence(address),
+                "agent_address": address,
+            }
+        }
+
+    def get_sequence(self, address: str) -> int:
+        query_msg = {"query_sequence": {"agent_address": address}}
+        sequence = self.query(query_msg)["sequence"]
+
+        return sequence
+
+
+class ServiceContract(LedgerContract):
+    def is_name_available(self, name: str):
+        query_msg = {"domain_record": {"domain": f"{name}.agent"}}
+        return self.query(query_msg)["is_available"]
+
+    def is_owner(self, name: str, wallet_address: str):
+        query_msg = {
+            "permissions": {
+                "domain": f"{name}.agent",
+                "owner": {"address": {"address": wallet_address}},
+            }
+        }
+        permission = self.query(query_msg)["permissions"]
+        if permission == "admin":
+            return True
+        return False
+
+    @staticmethod
+    def get_ownership_msg(name: str, wallet_address: str):
+        return {
+            "update_ownership": {
+                "domain": f"{name}.agent",
+                "owner": {"address": {"address": wallet_address}},
+                "permissions": "admin",
+            }
+        }
+
+    @staticmethod
+    def get_registration_msg(name: str, address: str):
+        return {
+            "register": {
+                "domain": f"{name}.agent",
+                "agent_address": address,
+            }
+        }
 
 
 if AGENT_NETWORK == AgentNetwork.FETCHAI_TESTNET:
@@ -24,7 +118,9 @@
 else:
     raise NotImplementedError
 
-_contract = LedgerContract(None, _ledger, CONTRACT_ALMANAC)
+
+_almanac_contract = AlmanacContract(None, _ledger, CONTRACT_ALMANAC)
+_service_contract = ServiceContract(None, _ledger, CONTRACT_SERVICE)
 
 
 def get_ledger() -> LedgerClient:
@@ -35,8 +131,12 @@ def get_faucet() -> FaucetApi:
     return _faucet_api
 
 
-def get_reg_contract() -> LedgerContract:
-    return _contract
+def get_almanac_contract() -> LedgerContract:
+    return _almanac_contract
+
+
+def get_service_contract() -> LedgerContract:
+    return _service_contract
 
 
 async def wait_for_tx_to_complete(

src/uagents/resolver.py

@@ -2,11 +2,11 @@
 from typing import Dict, Optional
 import random
 
-from uagents.network import get_reg_contract
+from uagents.network import get_almanac_contract
 
 
 def _query_record(agent_address: str, service: str) -> dict:
-    contract = get_reg_contract()
+    contract = get_almanac_contract()
     query_msg = {
         "query_record": {"agent_address": agent_address, "record_type": service}
     }