chore: set timeout to envelope expiry

python/src/uagents/asgi.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from logging import Logger
 from typing import (
     Any,
@@ -361,12 +361,17 @@ async def __call__(self, scope, receive, send):  #  pylint: disable=too-many-bra
 
         # wait for any queries to be resolved
         if expects_response:
-            response_msg, schema_digest = await asyncio.wait_for(
-                self._queries[env.sender], DEFAULT_ENVELOPE_TIMEOUT_SECONDS
+            timeout = (
+                env.expires - datetime.now(timezone.utc).timestamp()
+                if env.expires
+                else None
             )
-            if (env.expires is not None) and (
-                datetime.now() > datetime.fromtimestamp(env.expires)
-            ):
+            try:
+                response_msg, schema_digest = await asyncio.wait_for(
+                    self._queries[env.sender],
+                    timeout or DEFAULT_ENVELOPE_TIMEOUT_SECONDS,
+                )
+            except asyncio.TimeoutError:
                 response_msg = ErrorMessage(
                     error="Query envelope expired"
                 ).model_dump_json()

python/src/uagents/communication.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 import uuid
-from time import time
+from datetime import datetime, timezone
 from typing import List, Optional, Tuple, Type, Union
 
 import aiohttp
@@ -222,7 +222,7 @@ async def send_message_raw(
         target=destination_address,
         session=uuid.uuid4(),
         schema_digest=message_schema_digest,
-        expires=int(time()) + timeout,
+        expires=int(datetime.now(timezone.utc).timestamp()) + timeout,
     )
     env.encode_payload(message_body)
     if not is_user_address(sender_address) and isinstance(sender, Identity):

python/src/uagents/context.py

@@ -623,7 +623,7 @@ async def send(
         message: Model,
         sync: bool = False,
         timeout: int = DEFAULT_ENVELOPE_TIMEOUT_SECONDS,
-    ) -> MsgStatus:
+    ) -> Union[MsgStatus, Envelope]:
         """
         Send a message to the specified destination.