133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
import time
|
|
import uuid
|
|
from asyncio import Lock
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from api.domain.task_status import TaskStatus
|
|
|
|
|
|
@dataclass
|
|
class TaskRecord:
|
|
task_id: str
|
|
task: str
|
|
timeout: int
|
|
metadata: dict[str, Any] | None
|
|
status: TaskStatus = TaskStatus.queued
|
|
create_at: float = field(default_factory=time.time)
|
|
started_at: float | None = None
|
|
finished_at: float | None = None
|
|
result: str | None = None
|
|
error: str | None = None
|
|
raw_response: dict[str, Any] | None = None
|
|
human_intervention: dict[str, Any] | None = None
|
|
session_id: str | None = None
|
|
resume_token: str | None = None
|
|
awaiting_deadline: float | None = None
|
|
|
|
@property
|
|
def execution_time(self) -> float:
|
|
if self.started_at is None:
|
|
return 0
|
|
end = self.finished_at if self.finished_at is not None else time.time()
|
|
return max(0.0, end - self.started_at)
|
|
|
|
|
|
class TaskStore:
|
|
def __init__(self) -> None:
|
|
self._lock = Lock()
|
|
self._tasks: dict[str, TaskRecord] = {}
|
|
|
|
async def create(self, task: str, timeout: int, metadata: dict[str, Any] | None) -> TaskRecord:
|
|
task_id = uuid.uuid4().hex
|
|
rec = TaskRecord(task_id=task_id, task=task, timeout=timeout, metadata=metadata)
|
|
async with self._lock:
|
|
self._tasks[task_id] = rec
|
|
return rec
|
|
|
|
async def get(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
return self._tasks.get(task_id)
|
|
|
|
async def set_running(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.status = TaskStatus.running
|
|
rec.finished_at = None
|
|
rec.error = None
|
|
rec.human_intervention = None
|
|
rec.awaiting_deadline = None
|
|
if rec.started_at is None:
|
|
rec.started_at = time.time()
|
|
return rec
|
|
|
|
async def set_awaiting_captcha(
|
|
self,
|
|
task_id: str,
|
|
raw_response: dict[str, Any],
|
|
max_wait_seconds: int,
|
|
) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
payload = raw_response.get("human_intervention") or {}
|
|
rec.status = TaskStatus.awaiting_user_captcha
|
|
rec.raw_response = raw_response
|
|
rec.human_intervention = payload
|
|
rec.session_id = payload.get("session_id") or rec.session_id
|
|
rec.resume_token = payload.get("resume_token") or rec.resume_token
|
|
rec.awaiting_deadline = time.time() + max(1, int(max_wait_seconds))
|
|
rec.error = None
|
|
rec.finished_at = None
|
|
return rec
|
|
|
|
async def set_done(
|
|
self,
|
|
task_id: str,
|
|
success: bool,
|
|
raw_response: dict[str, Any] | None,
|
|
error: str | None,
|
|
result: str | None = None,
|
|
) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.finished_at = time.time()
|
|
rec.raw_response = raw_response
|
|
rec.error = error if error is not None else (
|
|
raw_response.get("error") if isinstance(raw_response, dict) else None
|
|
)
|
|
rec.result = result if result is not None else (
|
|
raw_response.get("result") if isinstance(raw_response, dict) else None
|
|
)
|
|
rec.human_intervention = None
|
|
rec.awaiting_deadline = None
|
|
rec.status = TaskStatus.succeeded if success else TaskStatus.failed
|
|
return rec
|
|
|
|
async def expire_if_needed(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.status != TaskStatus.awaiting_user_captcha:
|
|
return rec
|
|
if rec.awaiting_deadline is None or rec.awaiting_deadline > time.time():
|
|
return rec
|
|
|
|
rec.status = TaskStatus.failed
|
|
rec.finished_at = time.time()
|
|
rec.error = "CAPTCHA wait expired before the user completed verification."
|
|
rec.raw_response = {
|
|
"success": False,
|
|
"status": TaskStatus.failed.value,
|
|
"error": rec.error,
|
|
"error_code": "captcha_wait_expired",
|
|
}
|
|
rec.human_intervention = None
|
|
rec.awaiting_deadline = None
|
|
return rec
|