Retry logic when granting authz to new tasks

agentex/src/domain/use_cases/agents_acp_use_case.py

@@ -1,9 +1,14 @@
+import asyncio
 import json
+import random
 from collections.abc import AsyncIterator, Callable
 from typing import Annotated, Any
 
 from fastapi import Depends
 
+from src.adapters.authentication.exceptions import (
+    AuthenticationServiceUnavailableError,
+)
 from src.adapters.crud_store.exceptions import ItemDoesNotExist
 from src.api.schemas.authorization_types import (
     AgentexResource,
@@ -227,6 +232,30 @@ async def _execute_with_error_handling(
             await self.task_service.fail_task(task, str(e))
             raise e
 
+    async def grant_with_retry(self, task: TaskEntity, attempts: int = 0) -> None:
+        """Grant authorization for a task with retry"""
+        try:
+            await self.authorization_service.grant(
+                resource=AgentexResource.task(task.id),
+            )
+        except AuthenticationServiceUnavailableError as e:
+            if attempts < 3:
+                delay = 0.2 * (2**attempts) + random.uniform(0, 0.1)
+                logger.error(
+                    f"Authentication service unavailable: {e}. Retrying in {delay}s..."
+                )
+                await asyncio.sleep(delay)
+                return await self.grant_with_retry(task, attempts + 1)
+            else:
+                logger.error(
+                    f"Authentication service unavailable: {e}. Max retries reached."
+                )
+                raise e from e
+        except Exception as e:
+            logger.error(f"Error granting authorization for task {task.id}: {e}")
+            await self.task_service.fail_task(task, str(e))
+            raise e from e
+
     async def _get_or_create_task(
         self,
         *,
@@ -273,9 +302,7 @@ async def _get_or_create_task(
             agent=agent, task_name=task_name, task_params=task_params
         )
         logger.info(f"[agent_id={agent.id}] Created task {task.id}")
-        await self.authorization_service.grant(
-            resource=AgentexResource.task(task.id),
-        )
+        await self.grant_with_retry(task)
         return task
 
     async def handle_rpc_request(