add new tool: to_captcha

This commit is contained in:
VladislavIlin7 2026-05-17 01:03:55 +03:00
parent 8f86dbbdac
commit 4852345bf6
12 changed files with 716 additions and 35 deletions

View file

@ -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: