add new tool: to_captcha
This commit is contained in:
parent
50589232d6
commit
f1f32d8366
14 changed files with 1008 additions and 130 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue