feat: better exception handling (#133)

python/src/uagents/agent.py

@@ -2,6 +2,7 @@
 import functools
 from typing import Dict, List, Optional, Set, Union, Type, Tuple, Any, Coroutine
 import uuid
+from pydantic import ValidationError
 import requests
 
 from cosmpy.aerial.wallet import LocalWallet, PrivateKey
@@ -41,10 +42,12 @@ async def _run_interval(func: IntervalCallback, ctx: Context, period: float):
     while True:
         try:
             await func(ctx)
-        except OSError:
-            ctx.logger.exception("OS Error in interval handler")
-        except RuntimeError:
-            ctx.logger.exception("Runtime Error in interval handler")
+        except OSError as ex:
+            ctx.logger.exception(f"OS Error in interval handler: {ex}")
+        except RuntimeError as ex:
+            ctx.logger.exception(f"Runtime Error in interval handler: {ex}")
+        except Exception as ex:
+            ctx.logger.exception(f"Exception in interval handler: {ex}")
 
         await asyncio.sleep(period)
 
@@ -386,11 +389,25 @@ async def handle_message(
     async def _startup(self):
         await self._registration_loop()
         for handler in self._on_startup:
-            await handler(self._ctx)
+            try:
+                await handler(self._ctx)
+            except OSError as ex:
+                self._logger.exception(f"OS Error in startup handler: {ex}")
+            except RuntimeError as ex:
+                self._logger.exception(f"Runtime Error in startup handler: {ex}")
+            except Exception as ex:
+                self._logger.exception(f"Exception in startup handler: {ex}")
 
     async def _shutdown(self):
         for handler in self._on_shutdown:
-            await handler(self._ctx)
+            try:
+                await handler(self._ctx)
+            except OSError as ex:
+                self._logger.exception(f"OS Error in shutdown handler: {ex}")
+            except RuntimeError as ex:
+                self._logger.exception(f"Runtime Error in shutdown handler: {ex}")
+            except Exception as ex:
+                self._logger.exception(f"Exception in shutdown handler: {ex}")
 
     def setup(self):
         # register the internal agent protocol
@@ -436,7 +453,11 @@ async def _process_message_queue(self):
                 continue
 
             # parse the received message
-            recovered = model_class.parse_raw(message)
+            try:
+                recovered = model_class.parse_raw(message)
+            except ValidationError as ex:
+                self._logger.warning(f"Unable to parse message: {ex}")
+                continue
 
             context = Context(
                 self._identity.address,
@@ -475,7 +496,14 @@ async def _process_message_queue(self):
                     continue
 
             if handler is not None:
-                await handler(context, sender, recovered)
+                try:
+                    await handler(context, sender, recovered)
+                except OSError as ex:
+                    self._logger.exception(f"OS Error in message handler: {ex}")
+                except RuntimeError as ex:
+                    self._logger.exception(f"Runtime Error in message handler: {ex}")
+                except Exception as ex:
+                    self._logger.exception(f"Exception in message handler: {ex}")
 
 
 class Bureau:

python/src/uagents/context.py

@@ -242,13 +242,19 @@ async def send_raw(
         env.encode_payload(json_message)
         env.sign(self._identity)
 
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                endpoint, headers={"content-type": "application/json"}, data=env.json()
-            ) as resp:
-                success = resp.status == 200
-
-        if not success:
-            self._logger.exception(
-                f"Unable to send envelope to {destination_address} @ {endpoint}"
-            )
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    endpoint,
+                    headers={"content-type": "application/json"},
+                    data=env.json(),
+                ) as resp:
+                    success = resp.status == 200
+                if not success:
+                    self._logger.exception(
+                        f"Unable to send envelope to {destination_address} @ {endpoint}"
+                    )
+        except aiohttp.ClientConnectorError as ex:
+            self._logger.exception(f"Failed to connect to {endpoint}: {ex}")
+        except Exception as ex:
+            self._logger.exception(f"Failed to send message to {destination}: {ex}")

python/src/uagents/mailbox.py

@@ -84,17 +84,26 @@ async def _handle_envelope(self, payload: dict):
     async def process_deletion_queue(self):
         async with aiohttp.ClientSession() as session:
             while True:
-                env_payload = await self._envelopes_to_delete.get()
-                env_url = f"{self.http_prefix}://{self.base_url}/v1/mailbox/{env_payload['uuid']}"
-                self._logger.debug(f"Deleting message: {env_payload}")
-                async with session.delete(
-                    env_url,
-                    headers={"Authorization": f"token {self._access_token}"},
-                ) as resp:
-                    if resp.status != 200:
-                        self._logger.exception(
-                            f"Failed to delete envelope from inbox: {(await resp.text())}"
-                        )
+                try:
+                    env = await self._envelopes_to_delete.get()
+                    env_url = (
+                        f"{self.http_prefix}://{self.base_url}/v1/mailbox/{env['uuid']}"
+                    )
+                    self._logger.debug(f"Deleting message: {env}")
+                    async with session.delete(
+                        env_url,
+                        headers={"Authorization": f"token {self._access_token}"},
+                    ) as resp:
+                        if resp.status != 200:
+                            self._logger.exception(
+                                f"Failed to delete envelope from inbox: {(await resp.text())}"
+                            )
+                except ClientConnectorError as ex:
+                    self._logger.warning(f"Failed to connect to mailbox server: {ex}")
+                except Exception as ex:
+                    self._logger.exception(
+                        f"Got exception while processing deletion queue: {ex}"
+                    )
 
     async def _poll_server(self):
         async with aiohttp.ClientSession() as session:
@@ -133,10 +142,12 @@ async def _open_websocket_connection(self):
                         await self._handle_envelope(msg["payload"])
 
         except websockets.exceptions.ConnectionClosedError:
-            pass
+            self._logger.warning("Mailbox connection closed")
+            self._access_token = None
 
         except ConnectionRefusedError:
-            pass
+            self._logger.warning("Mailbox connection refused")
+            self._access_token = None
 
     async def _get_access_token(self):
         async with aiohttp.ClientSession() as session: