diff --git a/python/src/uagents/asgi.py b/python/src/uagents/asgi.py index 9867b79a9..17c2b6513 100644 --- 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 index 042f730a7..76409eb15 100644 --- 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 diff --git a/python/tests/test_server.py b/python/tests/test_server.py index a58fb0634..d9fa78ead 100644 --- a/python/tests/test_server.py +++ b/python/tests/test_server.py @@ -6,7 +6,8 @@ from uagents import Agent, Model from uagents.envelope import Envelope -from uagents.crypto import generate_user_address +from uagents.config import RESPONSE_TIME_HINT_SECONDS +from uagents.crypto import generate_user_address, Identity from uagents.query import enclose_response @@ -46,6 +47,7 @@ async def test_message_success(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -89,6 +91,7 @@ async def test_message_success_unsigned(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -134,6 +137,7 @@ async def test_message_success_sync(self): self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": { b"content-type": b"application/json", @@ -186,6 +190,7 @@ async def test_message_success_sync_unsigned(self): self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": { b"content-type": b"application/json", @@ -237,6 +242,7 @@ async def test_message_fail_wrong_path(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/bad/path", "headers": {b"content-type": b"application/json"}, }, @@ -279,6 +285,7 @@ async def test_message_fail_wrong_headers(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/badapp"}, }, @@ -311,6 +318,7 @@ async def test_message_fail_bad_data(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -352,6 +360,7 @@ async def test_message_fail_unsigned(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -394,6 +403,7 @@ async def test_message_fail_verify(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -436,6 +446,7 @@ async def test_message_fail_dispatch(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"content-type": b"application/json"}, }, @@ -465,6 +476,7 @@ async def test_request_fail_missing_header(self): await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {}, }, @@ -489,11 +501,76 @@ async def test_request_fail_missing_header(self): ] ) + async def test_request_fail_no_contents(self): + mock_send = AsyncMock() + with patch("uagents.asgi._read_asgi_body") as mock_receive: + mock_receive.return_value = None + await self.agent._server( + scope={ + "type": "http", + "method": "POST", + "path": "/submit", + "headers": {b"content-type": b"application/json"}, + }, + receive=None, + send=mock_send, + ) + mock_send.assert_has_calls( + [ + call( + { + "type": "http.response.start", + "status": 400, + "headers": [[b"content-type", b"application/json"]], + } + ), + call( + { + "type": "http.response.body", + "body": b'{"error": "empty or invalid payload"}', + } + ), + ] + ) + + async def test_request_fail_invalid_json(self): + mock_send = AsyncMock() + with patch("uagents.asgi._read_asgi_body") as mock_receive: + mock_receive.return_value = '{"bad", "json"}'.encode() + await self.agent._server( + scope={ + "type": "http", + "method": "POST", + "path": "/submit", + "headers": {b"content-type": b"application/json"}, + }, + receive=None, + send=mock_send, + ) + mock_send.assert_has_calls( + [ + call( + { + "type": "http.response.start", + "status": 400, + "headers": [[b"content-type", b"application/json"]], + } + ), + call( + { + "type": "http.response.body", + "body": b'{"error": "empty or invalid payload"}', + } + ), + ] + ) + async def test_request_from_browser(self): mock_send = AsyncMock() await self.agent._server( scope={ "type": "http", + "method": "POST", "path": "/submit", "headers": {b"User-Agent": b"Mozilla/5.0"}, }, @@ -518,6 +595,88 @@ async def test_request_from_browser(self): ] ) + async def test_head_no_address_header(self): + mock_send = AsyncMock() + await self.agent._server( + scope={ + "type": "http", + "method": "HEAD", + "path": "/submit", + "headers": {}, + }, + receive=None, + send=mock_send, + ) + mock_send.assert_has_calls( + [ + call( + { + "type": "http.response.start", + "status": 200, + "headers": [ + [b"x-uagents-status", b"indeterminate"], + ], + } + ), + ] + ) + + async def test_head_agent_ready(self): + mock_send = AsyncMock() + await self.agent._server( + scope={ + "type": "http", + "method": "HEAD", + "path": "/submit", + "headers": {b"x-uagents-address": self.agent.address.encode()}, + }, + receive=None, + send=mock_send, + ) + mock_send.assert_has_calls( + [ + call( + { + "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(), + ], + ], + } + ), + ] + ) + + async def test_head_agent_not_ready(self): + mock_send = AsyncMock() + await self.agent._server( + scope={ + "type": "http", + "method": "HEAD", + "path": "/submit", + "headers": {b"x-uagents-address": Identity.generate().address.encode()}, + }, + receive=None, + send=mock_send, + ) + mock_send.assert_has_calls( + [ + call( + { + "type": "http.response.start", + "status": 200, + "headers": [ + [b"x-uagents-status", b"not-ready"], + ], + } + ), + ] + ) + if __name__ == "__main__": unittest.main()