Spaces:
Running
Running
| import asyncio | |
| import os | |
| import time | |
| from fasta2a.task_manager import TaskManager | |
| from fasta2a.schema import SendMessageResponse | |
| try: | |
| from fasta2a.schema import SendMessageResult | |
| except ImportError: | |
| def SendMessageResult(**kwargs): | |
| return kwargs | |
| class BlockingTaskManager(TaskManager): | |
| def __init__(self, broker, storage, **kwargs): | |
| super().__init__(broker, storage, **kwargs) | |
| self.storage = storage | |
| async def send_message(self, request): | |
| response = await super().send_message(request) | |
| config = request["params"].get("configuration") or {} | |
| if not config.get("blocking", False): | |
| return response | |
| task = response["result"] | |
| task_id = task["id"] | |
| timeout = min( | |
| max(float(os.getenv("A2A_BLOCKING_TIMEOUT", "60")), 1), | |
| 90, | |
| ) | |
| deadline = time.monotonic() + timeout | |
| terminal_states = { | |
| "completed", | |
| "failed", | |
| "canceled", | |
| "rejected", | |
| "input-required", | |
| "auth-required", | |
| } | |
| while True: | |
| task = await self.storage.load_task(task_id) | |
| if task and task["status"]["state"] in terminal_states: | |
| return SendMessageResponse( | |
| jsonrpc="2.0", | |
| id=request["id"], | |
| result=task, | |
| ) | |
| if time.monotonic() >= deadline: | |
| return SendMessageResponse( | |
| jsonrpc="2.0", | |
| id=request["id"], | |
| result=task or response["result"], | |
| ) | |
| await asyncio.sleep(0.5) | |