add new tool: to_captcha

This commit is contained in:
VladislavIlin7 2026-04-21 23:32:09 +03:00
parent 50589232d6
commit f1f32d8366
14 changed files with 1008 additions and 130 deletions

View file

@ -1,21 +1,24 @@
import asyncio
from api.clients.browser_rpc_contracts import BrowserRpcError, BrowserRpcRunner
from api.domain.task_status import TaskStatus
from api.repositories.task_store import TaskRecord, TaskStore
class TaskService:
def __init__(
self,
store: TaskStore,
rpc_client: BrowserRpcRunner,
max_concurrency: int,
rpc_timeout_cap: float | None = None,
self,
store: TaskStore,
rpc_client: BrowserRpcRunner,
max_concurrency: int,
rpc_timeout_cap: float | None = None,
captcha_wait_timeout: int = 900,
) -> None:
self._store = store
self._rpc_client = rpc_client
self._semaphore = asyncio.Semaphore(max_concurrency)
self._rpc_timeout_cap = rpc_timeout_cap
self._captcha_wait_timeout = captcha_wait_timeout
self._background_tasks: set[asyncio.Task[None]] = set()
async def submit_task(self, task: str, timeout: int, metadata: dict | None) -> TaskRecord:
@ -26,7 +29,22 @@ class TaskService:
return record
async def get_task(self, task_id: str) -> TaskRecord | None:
return await self._store.get(task_id)
before = await self._store.get(task_id)
was_awaiting = bool(before is not None and before.status == TaskStatus.awaiting_user_captcha)
await self._store.expire_if_needed(task_id)
after = await self._store.get(task_id)
if (
was_awaiting
and after is not None
and after.status == TaskStatus.failed
and after.error
and "expired" in after.error.lower()
):
try:
await self._rpc_client.abort(task_id, reason=after.error)
except BrowserRpcError:
pass
return after
async def close(self) -> None:
if not self._background_tasks:
@ -37,6 +55,86 @@ class TaskService:
await asyncio.gather(*self._background_tasks, return_exceptions=True)
self._background_tasks.clear()
async def resume_captcha(self, task_id: str, user_response: str | None = None) -> TaskRecord | None:
rec = await self.get_task(task_id)
if rec is None:
return None
if rec.status == TaskStatus.failed:
return rec
if rec.status != TaskStatus.awaiting_user_captcha:
return rec
verify_raw = await self._rpc_client.verify_captcha(task_id)
if not verify_raw.get("verified"):
await self._store.set_awaiting_captcha(
task_id=task_id,
raw_response={
"success": False,
"status": TaskStatus.awaiting_user_captcha.value,
"error": "CAPTCHA is still present.",
"human_intervention": rec.human_intervention,
"verification": verify_raw,
"user_response": user_response,
},
max_wait_seconds=(rec.human_intervention or {}).get("verification", {}).get(
"max_wait_seconds", self._captcha_wait_timeout
),
)
return await self._store.get(task_id)
await self._store.set_running(task_id)
try:
rpc_timeout = float(rec.timeout)
if self._rpc_timeout_cap is not None:
rpc_timeout = min(rpc_timeout, self._rpc_timeout_cap)
raw = await asyncio.wait_for(
self._rpc_client.resume(task_id=task_id, timeout_sec=rpc_timeout),
timeout=float(rec.timeout) + 5,
)
except asyncio.TimeoutError:
await self._store.set_done(
task_id=task_id,
success=False,
raw_response=None,
error="Timeout exceeded after CAPTCHA resume.",
)
return await self._store.get(task_id)
except BrowserRpcError as exc:
await self._store.set_done(
task_id=task_id,
success=False,
raw_response=None,
error=str(exc),
)
return await self._store.get(task_id)
await self._apply_rpc_result(task_id, raw)
return await self._store.get(task_id)
async def abort_captcha(self, task_id: str, reason: str | None = None) -> TaskRecord | None:
rec = await self.get_task(task_id)
if rec is None:
return None
if rec.status == TaskStatus.awaiting_user_captcha:
try:
await self._rpc_client.abort(task_id, reason=reason)
except BrowserRpcError:
pass
await self._store.set_done(
task_id=task_id,
success=False,
raw_response={
"success": False,
"status": TaskStatus.failed.value,
"error_code": "captcha_aborted",
"reason": reason,
},
error=reason or "User aborted CAPTCHA flow.",
)
return await self._store.get(task_id)
async def _worker(self, task_id: str) -> None:
rec = await self._store.set_running(task_id)
if rec is None:
@ -49,17 +147,10 @@ class TaskService:
rpc_timeout = min(rpc_timeout, self._rpc_timeout_cap)
raw = await asyncio.wait_for(
self._rpc_client.run(task=rec.task, timeout_sec=rpc_timeout),
self._rpc_client.run(task_id=task_id, task=rec.task, timeout_sec=rpc_timeout),
timeout=float(rec.timeout) + 5,
)
success = bool(raw.get("success"))
await self._store.set_done(
task_id=task_id,
success=success,
raw_response=raw,
error=None,
result=raw.get("result") if isinstance(raw, dict) else None,
)
await self._apply_rpc_result(task_id, raw)
except asyncio.TimeoutError:
await self._store.set_done(
task_id=task_id,
@ -81,3 +172,22 @@ class TaskService:
raw_response=None,
error=f"Internal error: {exc}",
)
async def _apply_rpc_result(self, task_id: str, raw: dict | None) -> None:
raw = raw or {}
status = raw.get("status")
if status == TaskStatus.awaiting_user_captcha.value:
human = raw.get("human_intervention") or {}
verification = human.get("verification") or {}
max_wait_seconds = verification.get("max_wait_seconds", self._captcha_wait_timeout)
await self._store.set_awaiting_captcha(task_id, raw_response=raw, max_wait_seconds=max_wait_seconds)
return
success = bool(raw.get("success"))
await self._store.set_done(
task_id=task_id,
success=success,
raw_response=raw,
error=None,
result=raw.get("result") if isinstance(raw, dict) else None,
)