add new tool: to_captcha
This commit is contained in:
parent
8f86dbbdac
commit
4852345bf6
12 changed files with 716 additions and 35 deletions
|
|
@ -4,6 +4,7 @@ from asyncio import Event, Lock, Queue
|
|||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from api.domain.captcha_state import CaptchaState
|
||||
from api.domain.task_status import TaskStatus
|
||||
|
||||
|
||||
|
|
@ -25,6 +26,16 @@ class TaskRecord:
|
|||
cancel_requested: bool = False
|
||||
done_event: Event = field(default_factory=Event)
|
||||
|
||||
captcha_state: CaptchaState = CaptchaState.none
|
||||
captcha_kind: str | None = None
|
||||
captcha_reason: str | None = None
|
||||
captcha_view_url: str | None = None
|
||||
captcha_notified_at: float | None = None
|
||||
captcha_solved_at: float | None = None
|
||||
captcha_deadline: float | None = None
|
||||
captcha_extra_seconds: int = 0
|
||||
captcha_event: Event = field(default_factory=Event)
|
||||
|
||||
@property
|
||||
def execution_time(self) -> float:
|
||||
if self.started_at is None:
|
||||
|
|
@ -40,13 +51,7 @@ class TaskStore:
|
|||
self._thread_index: dict[str, list[str]] = {}
|
||||
self._subscribers: dict[str, set[Queue[dict[str, Any]]]] = {}
|
||||
|
||||
async def create(
|
||||
self,
|
||||
task: str,
|
||||
timeout: int,
|
||||
metadata: dict[str, Any] | None,
|
||||
thread_id: str = "default",
|
||||
) -> TaskRecord:
|
||||
async def create(self, task: str, timeout: int, metadata: dict[str, Any] | None, thread_id: str = "default") -> TaskRecord:
|
||||
task_id = uuid.uuid4().hex
|
||||
rec = TaskRecord(task_id=task_id, thread_id=thread_id, task=task, timeout=timeout, metadata=metadata)
|
||||
async with self._lock:
|
||||
|
|
@ -75,25 +80,15 @@ class TaskStore:
|
|||
rec.started_at = time.time()
|
||||
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,
|
||||
history: list[dict[str, Any]] | None = None,
|
||||
) -> TaskRecord | None:
|
||||
async def set_done(self, task_id: str, success: bool, raw_response: dict[str, Any] | None, error: str | None, result: str | None = None, history: list[dict[str, Any]] | 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.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.history = list(history or [])
|
||||
rec.status = TaskStatus.succeeded if success else TaskStatus.failed
|
||||
rec.done_event.set()
|
||||
|
|
@ -132,7 +127,6 @@ class TaskStore:
|
|||
return False, False
|
||||
if rec.status in (TaskStatus.queued, TaskStatus.running):
|
||||
return True, False
|
||||
|
||||
del self._tasks[task_id]
|
||||
thread_list = self._thread_index.get(rec.thread_id, [])
|
||||
if task_id in thread_list:
|
||||
|
|
@ -140,6 +134,93 @@ class TaskStore:
|
|||
self._subscribers.pop(task_id, None)
|
||||
return True, True
|
||||
|
||||
async def set_captcha_awaiting(self, task_id: str, kind: str | None, reason: str | None, view_url: str | None, timeout_seconds: int) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
now = time.time()
|
||||
rec.captcha_state = CaptchaState.awaiting
|
||||
rec.captcha_kind = kind
|
||||
rec.captcha_reason = reason
|
||||
rec.captcha_view_url = view_url
|
||||
rec.captcha_notified_at = now
|
||||
rec.captcha_solved_at = None
|
||||
rec.captcha_extra_seconds = 0
|
||||
rec.captcha_deadline = now + max(1, int(timeout_seconds))
|
||||
rec.captcha_event.clear()
|
||||
return rec
|
||||
|
||||
async def set_captcha_solved(self, task_id: str) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
if rec.captcha_state in (CaptchaState.none, CaptchaState.solved):
|
||||
return rec
|
||||
rec.captcha_state = CaptchaState.solved
|
||||
rec.captcha_solved_at = time.time()
|
||||
rec.captcha_event.set()
|
||||
return rec
|
||||
|
||||
async def set_captcha_timeout_prompt(self, task_id: str) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
if rec.captcha_state != CaptchaState.awaiting:
|
||||
return rec
|
||||
rec.captcha_state = CaptchaState.timeout_prompt
|
||||
rec.captcha_event.set()
|
||||
return rec
|
||||
|
||||
async def set_captcha_extended(self, task_id: str, extra_seconds: int) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
if rec.captcha_state not in (CaptchaState.timeout_prompt, CaptchaState.awaiting):
|
||||
return rec
|
||||
extra = max(1, int(extra_seconds))
|
||||
rec.captcha_extra_seconds += extra
|
||||
base = rec.captcha_deadline or time.time()
|
||||
rec.captcha_deadline = max(base, time.time()) + extra
|
||||
rec.captcha_state = CaptchaState.extended
|
||||
rec.captcha_event.set()
|
||||
return rec
|
||||
|
||||
async def set_captcha_aborted(self, task_id: str, reason: str | None = None) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
if rec.captcha_state in (CaptchaState.none, CaptchaState.aborted, CaptchaState.failed):
|
||||
return rec
|
||||
rec.captcha_state = CaptchaState.aborted
|
||||
rec.captcha_reason = reason or rec.captcha_reason
|
||||
rec.captcha_event.set()
|
||||
return rec
|
||||
|
||||
async def set_captcha_failed(self, task_id: str, error: str | None = None) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
rec.captcha_state = CaptchaState.failed
|
||||
if error:
|
||||
rec.captcha_reason = error
|
||||
rec.captcha_event.set()
|
||||
return rec
|
||||
|
||||
async def reset_captcha(self, task_id: str) -> TaskRecord | None:
|
||||
async with self._lock:
|
||||
rec = self._tasks.get(task_id)
|
||||
if rec is None:
|
||||
return None
|
||||
rec.captcha_state = CaptchaState.none
|
||||
rec.captcha_event.clear()
|
||||
return rec
|
||||
|
||||
async def subscribe(self, task_id: str) -> Queue[dict[str, Any]] | None:
|
||||
queue: Queue[dict[str, Any]] = Queue()
|
||||
async with self._lock:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue