feat: readiness probe (#194)

fetchai · Oct 25, 2023 · 516934b · 516934b
1 parent e361ead
commit 516934b
Show file tree

Hide file tree

Showing 3 changed files with 268 additions and 36 deletions.
diff --git a/python/src/uagents/asgi.py b/python/src/uagents/asgi.py
@@ -8,7 +8,7 @@
 import uvicorn
 from requests.structures import CaseInsensitiveDict
 
-from uagents.config import get_logger
+from uagents.config import get_logger, RESPONSE_TIME_HINT_SECONDS
 from uagents.crypto import is_user_address
 from uagents.dispatch import dispatcher
 from uagents.envelope import Envelope
@@ -70,6 +70,86 @@ def server(self):
         """
         return self._server
 
+    async def handle_readiness_probe(self, headers: CaseInsensitiveDict, send):
+        """
+        Handle a readiness probe sent via the HEAD method.
+        """
+        if b"x-uagents-address" not in headers:
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 200,
+                    "headers": [
+                        [b"x-uagents-status", b"indeterminate"],
+                    ],
+                }
+            )
+        else:
+            address = headers[b"x-uagents-address"].decode()
+            if not dispatcher.contains(address):
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 200,
+                        "headers": [
+                            [b"x-uagents-status", b"not-ready"],
+                        ],
+                    }
+                )
+            else:
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 200,
+                        "headers": [
+                            [b"x-uagents-status", b"ready"],
+                            [
+                                b"x-uagents-response-time-hint",
+                                str(RESPONSE_TIME_HINT_SECONDS).encode(),
+                            ],
+                        ],
+                    }
+                )
+        return
+
+    async def handle_missing_content_type(self, headers: CaseInsensitiveDict, send):
+        """
+        Handle missing content type header.
+        """
+        # if connecting from browser, return a 200 OK
+        if b"user-agent" in headers:
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 200,
+                    "headers": [
+                        [b"content-type", b"application/json"],
+                    ],
+                }
+            )
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": b'{"status": "OK - Agent is running"}',
+                }
+            )
+        else:  # otherwise, return a 400 Bad Request
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 400,
+                    "headers": [
+                        [b"content-type", b"application/json"],
+                    ],
+                }
+            )
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": b'{"error": "missing header: content-type"}',
+                }
+            )
+
     async def serve(self):
         """
         Start the server.
@@ -110,40 +190,13 @@ async def __call__(
 
         headers = CaseInsensitiveDict(scope.get("headers", {}))
 
+        request_method = scope["method"]
+        if request_method == "HEAD":
+            await self.handle_readiness_probe(headers, send)
+            return
+
         if b"content-type" not in headers:
-            # if connecting from browser, return a 200 OK
-            if b"user-agent" in headers:
-                await send(
-                    {
-                        "type": "http.response.start",
-                        "status": 200,
-                        "headers": [
-                            [b"content-type", b"application/json"],
-                        ],
-                    }
-                )
-                await send(
-                    {
-                        "type": "http.response.body",
-                        "body": b'{"status": "OK - Agent is running"}',
-                    }
-                )
-            else:  # otherwise, return a 400 Bad Request
-                await send(
-                    {
-                        "type": "http.response.start",
-                        "status": 400,
-                        "headers": [
-                            [b"content-type", b"application/json"],
-                        ],
-                    }
-                )
-                await send(
-                    {
-                        "type": "http.response.body",
-                        "body": b'{"error": "missing header: content-type"}',
-                    }
-                )
+            await self.handle_missing_content_type(headers, send)
             return
 
         if b"application/json" not in headers[b"content-type"]:
@@ -166,7 +219,26 @@ async def __call__(
 
         # read the entire payload
         raw_contents = await _read_asgi_body(receive)
-        contents = json.loads(raw_contents.decode())
+
+        try:
+            contents = json.loads(raw_contents.decode())
+        except (AttributeError, UnicodeDecodeError, json.JSONDecodeError):
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 400,
+                    "headers": [
+                        [b"content-type", b"application/json"],
+                    ],
+                }
+            )
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": b'{"error": "empty or invalid payload"}',
+                }
+            )
+            return
 
         try:
             env: Envelope = Envelope.parse_obj(contents)

diff --git a/python/src/uagents/config.py b/python/src/uagents/config.py
@@ -37,6 +37,7 @@
 ALMANAC_API_URL = AGENTVERSE_URL + "/v1/almanac/"
 MAILBOX_POLL_INTERVAL_SECONDS = 1.0
 
+RESPONSE_TIME_HINT_SECONDS = 5
 DEFAULT_ENVELOPE_TIMEOUT_SECONDS = 30
 DEFAULT_MAX_ENDPOINTS = 10
 DEFAULT_SEARCH_LIMIT = 100