|
| 1 | +import asyncio |
1 | 2 | import json |
| 3 | +import random |
2 | 4 | from collections.abc import AsyncIterator, Callable |
3 | 5 | from typing import Annotated, Any |
4 | 6 |
|
5 | 7 | from fastapi import Depends |
6 | 8 |
|
| 9 | +from src.adapters.authentication.exceptions import ( |
| 10 | + AuthenticationServiceUnavailableError, |
| 11 | +) |
7 | 12 | from src.adapters.crud_store.exceptions import ItemDoesNotExist |
8 | 13 | from src.api.schemas.authorization_types import ( |
9 | 14 | AgentexResource, |
@@ -227,6 +232,30 @@ async def _execute_with_error_handling( |
227 | 232 | await self.task_service.fail_task(task, str(e)) |
228 | 233 | raise e |
229 | 234 |
|
| 235 | + async def grant_with_retry(self, task: TaskEntity, attempts: int = 0) -> None: |
| 236 | + """Grant authorization for a task with retry""" |
| 237 | + try: |
| 238 | + await self.authorization_service.grant( |
| 239 | + resource=AgentexResource.task(task.id), |
| 240 | + ) |
| 241 | + except AuthenticationServiceUnavailableError as e: |
| 242 | + if attempts < 3: |
| 243 | + delay = 0.2 * (2**attempts) + random.uniform(0, 0.1) |
| 244 | + logger.error( |
| 245 | + f"Authentication service unavailable: {e}. Retrying in {delay}s..." |
| 246 | + ) |
| 247 | + await asyncio.sleep(delay) |
| 248 | + return await self.grant_with_retry(task, attempts + 1) |
| 249 | + else: |
| 250 | + logger.error( |
| 251 | + f"Authentication service unavailable: {e}. Max retries reached." |
| 252 | + ) |
| 253 | + raise e from e |
| 254 | + except Exception as e: |
| 255 | + logger.error(f"Error granting authorization for task {task.id}: {e}") |
| 256 | + await self.task_service.fail_task(task, str(e)) |
| 257 | + raise e from e |
| 258 | + |
230 | 259 | async def _get_or_create_task( |
231 | 260 | self, |
232 | 261 | *, |
@@ -273,9 +302,7 @@ async def _get_or_create_task( |
273 | 302 | agent=agent, task_name=task_name, task_params=task_params |
274 | 303 | ) |
275 | 304 | logger.info(f"[agent_id={agent.id}] Created task {task.id}") |
276 | | - await self.authorization_service.grant( |
277 | | - resource=AgentexResource.task(task.id), |
278 | | - ) |
| 305 | + await self.grant_with_retry(task) |
279 | 306 | return task |
280 | 307 |
|
281 | 308 | async def handle_rpc_request( |
|
0 commit comments