46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
from collections.abc import Awaitable, Callable
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.errors import AgentError, RecoverableAgentError
|
|
from app.repository.agent_run_repository import AgentRunRepository
|
|
|
|
|
|
class AgentRunWorker:
|
|
def __init__(
|
|
self,
|
|
session: AsyncSession,
|
|
execute_run: Callable[[str], Awaitable[None]],
|
|
worker_id: str,
|
|
lease_seconds: int = 60,
|
|
retry_limit: int = 3,
|
|
) -> None:
|
|
self.repository = AgentRunRepository(session)
|
|
self.execute_run = execute_run
|
|
self.worker_id = worker_id
|
|
self.lease_seconds = lease_seconds
|
|
self.retry_limit = retry_limit
|
|
|
|
async def claim_and_execute(self, run_id: str) -> bool:
|
|
async with self.repository.session.begin():
|
|
claimed = await self.repository.claim(run_id, self.worker_id, self.lease_seconds)
|
|
if not claimed:
|
|
return False
|
|
try:
|
|
await self.execute_run(run_id)
|
|
except RecoverableAgentError as exc:
|
|
async with self.repository.session.begin():
|
|
await self.repository.record_failure(run_id, exc.code, self.retry_limit)
|
|
return False
|
|
except AgentError as exc:
|
|
async with self.repository.session.begin():
|
|
await self.repository.record_failure(run_id, exc.code, 0)
|
|
return False
|
|
return True
|
|
|
|
async def renew(self, run_id: str) -> bool:
|
|
return await self.repository.renew(run_id, self.worker_id, self.lease_seconds)
|
|
|
|
async def cancellation_requested(self, run_id: str) -> bool:
|
|
return await self.repository.is_cancel_requested(run_id)
|