From 61aa96a0959cbdc0b7f342908fef2a54d92301a5 Mon Sep 17 00:00:00 2001
From: James Riehl <james.riehl@fetch.ai>
Date: Thu, 19 Dec 2024 14:19:08 +0000
Subject: [PATCH] fix(core): proxy endpoint flow

---
 python/src/uagents/agent.py   |  1 -
 python/src/uagents/mailbox.py | 26 ++++----------------------
 2 files changed, 4 insertions(+), 23 deletions(-)

diff --git a/python/src/uagents/agent.py b/python/src/uagents/agent.py
index c59e5769..ee76d93d 100644
--- a/python/src/uagents/agent.py
+++ b/python/src/uagents/agent.py
@@ -447,7 +447,6 @@ async def _handle_prove(_ctx: Context, request: AgentverseConnectRequest):
                 return await register_in_agentverse(
                     request,
                     self._identity,
-                    self._endpoints,
                     self._agentverse,
                     agent_details,
                 )
diff --git a/python/src/uagents/mailbox.py b/python/src/uagents/mailbox.py
index e8d5bb1c..fa1063f1 100644
--- a/python/src/uagents/mailbox.py
+++ b/python/src/uagents/mailbox.py
@@ -24,6 +24,7 @@
 class AgentverseConnectRequest(Model):
     user_token: str
     agent_type: AgentType
+    endpoint: Optional[str] = None
 
 
 class ChallengeRequest(BaseModel):
@@ -50,7 +51,7 @@ class RegistrationRequest(BaseModel):
     challenge: str
     challenge_response: str
     agent_type: AgentType
-    endpoints: Optional[list[AgentEndpoint]] = None
+    endpoint: Optional[str] = None
 
 
 class RegistrationResponse(Model):
@@ -83,22 +84,9 @@ def is_mailbox_agent(
     return any([f"{agentverse.url}/v1/submit" in ep.url for ep in endpoints])
 
 
-def is_proxy_agent(
-    endpoints: list[AgentEndpoint], agentverse: AgentverseConfig
-) -> bool:
-    """
-    Check if the agent is a proxy agent.
-
-    Returns:
-        bool: True if the agent is a proxy agent, False otherwise.
-    """
-    return any([f"{agentverse.url}/v1/proxy/submit" in ep.url for ep in endpoints])
-
-
 async def register_in_agentverse(
     request: AgentverseConnectRequest,
     identity: Identity,
-    endpoints: list[AgentEndpoint],
     agentverse: AgentverseConfig,
     agent_details: Optional[AgentUpdates] = None,
 ) -> RegistrationResponse:
@@ -108,8 +96,8 @@ async def register_in_agentverse(
     Args:
         request (AgentverseConnectRequest): Request object
         identity (Identity): Agent identity object
-        endpoints (list[AgentEndpoint]): Endpoints of the agent
         agentverse (AgentverseConfig): Agentverse configuration
+        agent_details (Optional[AgentUpdates]): Agent details (name, readme, avatar_url)
 
     Returns:
         RegistrationResponse: Registration
@@ -138,7 +126,7 @@ async def register_in_agentverse(
                 address=identity.address,
                 challenge=challenge.challenge,
                 challenge_response=identity.sign(challenge.challenge.encode()),
-                endpoints=endpoints,
+                endpoint=request.endpoint,
                 agent_type=request.agent_type,
             ).model_dump_json(),
             headers={
@@ -164,12 +152,6 @@ async def register_in_agentverse(
             request.user_token, identity.address, agent_details, agentverse
         )
 
-    if request.agent_type == "mailbox" and not is_mailbox_agent(endpoints, agentverse):
-        logger.exception(
-            f"Agent endpoints {endpoints} do not match registered agent type: {request.agent_type}"
-            f"Please restart agent with endpoint='{agentverse.url}/v1/submit'"
-        )
-
     return registration_response